burn_nn/modules/rnn/
gru.rs

1use burn_core as burn;
2
3use super::gate_controller::GateController;
4use burn::config::Config;
5use burn::module::Initializer;
6use burn::module::Module;
7use burn::module::{Content, DisplaySettings, ModuleDisplay};
8use burn::tensor::Tensor;
9use burn::tensor::activation;
10use burn::tensor::backend::Backend;
11
12/// Configuration to create a [gru](Gru) module using the [init function](GruConfig::init).
13#[derive(Config, Debug)]
14pub struct GruConfig {
15    /// The size of the input features.
16    pub d_input: usize,
17    /// The size of the hidden state.
18    pub d_hidden: usize,
19    /// If a bias should be applied during the Gru transformation.
20    pub bias: bool,
21    /// If reset gate should be applied after weight multiplication.
22    ///
23    /// This configuration option controls how the reset gate is applied to the hidden state.
24    /// * `true` - (Default) Match the initial arXiv version of the paper [Learning Phrase Representations using RNN Encoder-Decoder for
25    ///   Statistical Machine Translation (v1)](https://arxiv.org/abs/1406.1078v1) and apply the reset gate after multiplication by
26    ///   the weights. This matches the behavior of [PyTorch GRU](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html#torch.nn.GRU).
27    /// * `false` - Match the most recent revision of [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine
28    ///   Translation (v3)](https://arxiv.org/abs/1406.1078) and apply the reset gate before the weight multiplication.
29    ///
30    /// The differing implementations can give slightly different numerical results and have different efficiencies. For more
31    /// motivation for why the `true` can be more efficient see [Optimizing RNNs with Differentiable Graphs](https://svail.github.io/diff_graphs).
32    ///
33    /// To set this field to `false` use [`with_reset_after`](`GruConfig::with_reset_after`).
34    #[config(default = "true")]
35    pub reset_after: bool,
36    /// Gru initializer
37    #[config(default = "Initializer::XavierNormal{gain:1.0}")]
38    pub initializer: Initializer,
39}
40
41/// The Gru (Gated recurrent unit) module. This implementation is for a unidirectional, stateless, Gru.
42///
43/// Introduced in the paper: [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://arxiv.org/abs/1406.1078).
44///
45/// Should be created with [GruConfig].
46#[derive(Module, Debug)]
47#[module(custom_display)]
48pub struct Gru<B: Backend> {
49    /// The update gate controller.
50    pub update_gate: GateController<B>,
51    /// The reset gate controller.
52    pub reset_gate: GateController<B>,
53    /// The new gate controller.
54    pub new_gate: GateController<B>,
55    /// The size of the hidden state.
56    pub d_hidden: usize,
57    /// If reset gate should be applied after weight multiplication.
58    pub reset_after: bool,
59}
60
61impl<B: Backend> ModuleDisplay for Gru<B> {
62    fn custom_settings(&self) -> Option<DisplaySettings> {
63        DisplaySettings::new()
64            .with_new_line_after_attribute(false)
65            .optional()
66    }
67
68    fn custom_content(&self, content: Content) -> Option<Content> {
69        let [d_input, _] = self.update_gate.input_transform.weight.shape().dims();
70        let bias = self.update_gate.input_transform.bias.is_some();
71
72        content
73            .add("d_input", &d_input)
74            .add("d_hidden", &self.d_hidden)
75            .add("bias", &bias)
76            .add("reset_after", &self.reset_after)
77            .optional()
78    }
79}
80
81impl GruConfig {
82    /// Initialize a new [gru](Gru) module.
83    pub fn init<B: Backend>(&self, device: &B::Device) -> Gru<B> {
84        let d_output = self.d_hidden;
85
86        let update_gate = GateController::new(
87            self.d_input,
88            d_output,
89            self.bias,
90            self.initializer.clone(),
91            device,
92        );
93        let reset_gate = GateController::new(
94            self.d_input,
95            d_output,
96            self.bias,
97            self.initializer.clone(),
98            device,
99        );
100        let new_gate = GateController::new(
101            self.d_input,
102            d_output,
103            self.bias,
104            self.initializer.clone(),
105            device,
106        );
107
108        Gru {
109            update_gate,
110            reset_gate,
111            new_gate,
112            d_hidden: self.d_hidden,
113            reset_after: self.reset_after,
114        }
115    }
116}
117
118impl<B: Backend> Gru<B> {
119    /// Applies the forward pass on the input tensor. This GRU implementation
120    /// returns a state tensor with dimensions `[batch_size, sequence_length, hidden_size]`.
121    ///
122    /// # Parameters
123    /// - batched_input: `[batch_size, sequence_length, input_size]`.
124    /// - state: An optional tensor representing an initial cell state with dimensions
125    ///   `[batch_size, hidden_size]`. If none is provided, an empty state will be used.
126    ///
127    /// # Returns
128    /// - output: `[batch_size, sequence_length, hidden_size]`
129    pub fn forward(
130        &self,
131        batched_input: Tensor<B, 3>,
132        state: Option<Tensor<B, 2>>,
133    ) -> Tensor<B, 3> {
134        let device = batched_input.device();
135        let [batch_size, seq_length, _] = batched_input.shape().dims();
136
137        let mut batched_hidden_state =
138            Tensor::empty([batch_size, seq_length, self.d_hidden], &device);
139
140        let mut hidden_t = match state {
141            Some(state) => state,
142            None => Tensor::zeros([batch_size, self.d_hidden], &device),
143        };
144
145        for (t, input_t) in batched_input.iter_dim(1).enumerate() {
146            let input_t = input_t.squeeze_dim(1);
147            // u(pdate)g(ate) tensors
148            let biased_ug_input_sum =
149                self.gate_product(&input_t, &hidden_t, None, &self.update_gate);
150            let update_values = activation::sigmoid(biased_ug_input_sum); // Colloquially referred to as z(t)
151
152            // r(eset)g(ate) tensors
153            let biased_rg_input_sum =
154                self.gate_product(&input_t, &hidden_t, None, &self.reset_gate);
155            let reset_values = activation::sigmoid(biased_rg_input_sum); // Colloquially referred to as r(t)
156
157            // n(ew)g(ate) tensor
158            let biased_ng_input_sum = if self.reset_after {
159                self.gate_product(&input_t, &hidden_t, Some(&reset_values), &self.new_gate)
160            } else {
161                let reset_t = hidden_t.clone().mul(reset_values); // Passed as input to new_gate
162                self.gate_product(&input_t, &reset_t, None, &self.new_gate)
163            };
164            let candidate_state = biased_ng_input_sum.tanh(); // Colloquially referred to as g(t)
165
166            // calculate linear interpolation between previous hidden state and candidate state:
167            // g(t) * (1 - z(t)) + z(t) * hidden_t
168            hidden_t = candidate_state
169                .clone()
170                .mul(update_values.clone().sub_scalar(1).mul_scalar(-1)) // (1 - z(t)) = -(z(t) - 1)
171                + update_values.clone().mul(hidden_t);
172
173            let unsqueezed_hidden_state = hidden_t.clone().unsqueeze_dim(1);
174
175            batched_hidden_state = batched_hidden_state.slice_assign(
176                [0..batch_size, t..(t + 1), 0..self.d_hidden],
177                unsqueezed_hidden_state,
178            );
179        }
180
181        batched_hidden_state
182    }
183
184    /// Helper function for performing weighted matrix product for a gate and adds
185    /// bias, if any, and optionally applies reset to hidden state.
186    ///
187    ///  Mathematically, performs `Wx*X + r .* (Wh*H + b)`, where:
188    ///     Wx = weight matrix for the connection to input vector X
189    ///     Wh = weight matrix for the connection to hidden state H
190    ///     X = input vector
191    ///     H = hidden state
192    ///     b = bias terms
193    ///     r = reset state
194    fn gate_product(
195        &self,
196        input: &Tensor<B, 2>,
197        hidden: &Tensor<B, 2>,
198        reset: Option<&Tensor<B, 2>>,
199        gate: &GateController<B>,
200    ) -> Tensor<B, 2> {
201        let input_product = input.clone().matmul(gate.input_transform.weight.val());
202        let hidden_product = hidden.clone().matmul(gate.hidden_transform.weight.val());
203
204        let input_bias = gate
205            .input_transform
206            .bias
207            .as_ref()
208            .map(|bias_param| bias_param.val());
209        let hidden_bias = gate
210            .hidden_transform
211            .bias
212            .as_ref()
213            .map(|bias_param| bias_param.val());
214
215        match (input_bias, hidden_bias, reset) {
216            (Some(input_bias), Some(hidden_bias), Some(r)) => {
217                input_product
218                    + input_bias.unsqueeze()
219                    + r.clone().mul(hidden_product + hidden_bias.unsqueeze())
220            }
221            (Some(input_bias), Some(hidden_bias), None) => {
222                input_product + input_bias.unsqueeze() + hidden_product + hidden_bias.unsqueeze()
223            }
224            (Some(input_bias), None, Some(r)) => {
225                input_product + input_bias.unsqueeze() + r.clone().mul(hidden_product)
226            }
227            (Some(input_bias), None, None) => {
228                input_product + input_bias.unsqueeze() + hidden_product
229            }
230            (None, Some(hidden_bias), Some(r)) => {
231                input_product + r.clone().mul(hidden_product + hidden_bias.unsqueeze())
232            }
233            (None, Some(hidden_bias), None) => {
234                input_product + hidden_product + hidden_bias.unsqueeze()
235            }
236            (None, None, Some(r)) => input_product + r.clone().mul(hidden_product),
237            (None, None, None) => input_product + hidden_product,
238        }
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245    use crate::{LinearRecord, TestBackend};
246    use burn::module::Param;
247    use burn::tensor::{Distribution, TensorData};
248    use burn::tensor::{Tolerance, ops::FloatElem};
249
250    type FT = FloatElem<TestBackend>;
251
252    fn init_gru<B: Backend>(reset_after: bool, device: &B::Device) -> Gru<B> {
253        fn create_gate_controller<B: Backend>(
254            weights: f32,
255            biases: f32,
256            d_input: usize,
257            d_output: usize,
258            bias: bool,
259            initializer: Initializer,
260            device: &B::Device,
261        ) -> GateController<B> {
262            let record_1 = LinearRecord {
263                weight: Param::from_data(TensorData::from([[weights]]), device),
264                bias: Some(Param::from_data(TensorData::from([biases]), device)),
265            };
266            let record_2 = LinearRecord {
267                weight: Param::from_data(TensorData::from([[weights]]), device),
268                bias: Some(Param::from_data(TensorData::from([biases]), device)),
269            };
270            GateController::create_with_weights(
271                d_input,
272                d_output,
273                bias,
274                initializer,
275                record_1,
276                record_2,
277            )
278        }
279
280        let config = GruConfig::new(1, 1, false).with_reset_after(reset_after);
281        let mut gru = config.init::<B>(device);
282
283        gru.update_gate = create_gate_controller(
284            0.5,
285            0.0,
286            1,
287            1,
288            false,
289            Initializer::XavierNormal { gain: 1.0 },
290            device,
291        );
292        gru.reset_gate = create_gate_controller(
293            0.6,
294            0.0,
295            1,
296            1,
297            false,
298            Initializer::XavierNormal { gain: 1.0 },
299            device,
300        );
301        gru.new_gate = create_gate_controller(
302            0.7,
303            0.0,
304            1,
305            1,
306            false,
307            Initializer::XavierNormal { gain: 1.0 },
308            device,
309        );
310        gru
311    }
312
313    /// Test forward pass with simple input vector.
314    ///
315    /// z_t = sigmoid(0.5*0.1 + 0.5*0) = 0.5125
316    /// r_t = sigmoid(0.6*0.1 + 0.*0) = 0.5150
317    /// g_t = tanh(0.7*0.1 + 0.7*0) = 0.0699
318    ///
319    /// h_t = z_t * h' + (1 - z_t) * g_t = 0.0341
320    #[test]
321    fn tests_forward_single_input_single_feature() {
322        let device = Default::default();
323        TestBackend::seed(&device, 0);
324
325        let mut gru = init_gru::<TestBackend>(false, &device);
326
327        let input = Tensor::<TestBackend, 3>::from_data(TensorData::from([[[0.1]]]), &device);
328        let expected = TensorData::from([[0.034]]);
329
330        // Reset gate applied to hidden state before the matrix multiplication
331        let state = gru.forward(input.clone(), None);
332
333        let output = state
334            .select(0, Tensor::arange(0..1, &device))
335            .squeeze_dim::<2>(0);
336
337        let tolerance = Tolerance::default();
338        output
339            .to_data()
340            .assert_approx_eq::<FT>(&expected, tolerance);
341
342        // Reset gate applied to hidden state after the matrix multiplication
343        gru.reset_after = true; // override forward behavior
344        let state = gru.forward(input, None);
345
346        let output = state
347            .select(0, Tensor::arange(0..1, &device))
348            .squeeze_dim::<2>(0);
349
350        output
351            .to_data()
352            .assert_approx_eq::<FT>(&expected, tolerance);
353    }
354
355    #[test]
356    fn tests_forward_seq_len_3() {
357        let device = Default::default();
358        TestBackend::seed(&device, 0);
359        let mut gru = init_gru::<TestBackend>(true, &device);
360
361        let input =
362            Tensor::<TestBackend, 3>::from_data(TensorData::from([[[0.1], [0.2], [0.3]]]), &device);
363        let expected = TensorData::from([[0.0341], [0.0894], [0.1575]]);
364
365        let result = gru.forward(input.clone(), None);
366        let output = result
367            .select(0, Tensor::arange(0..1, &device))
368            .squeeze_dim::<2>(0);
369
370        let tolerance = Tolerance::default();
371        output
372            .to_data()
373            .assert_approx_eq::<FT>(&expected, tolerance);
374
375        // Reset gate applied to hidden state before the matrix multiplication
376        gru.reset_after = false; // override forward behavior
377        let state = gru.forward(input, None);
378
379        let output = state
380            .select(0, Tensor::arange(0..1, &device))
381            .squeeze_dim::<2>(0);
382
383        output
384            .to_data()
385            .assert_approx_eq::<FT>(&expected, tolerance);
386    }
387
388    #[test]
389    fn test_batched_forward_pass() {
390        let device = Default::default();
391        let gru = GruConfig::new(64, 1024, true).init::<TestBackend>(&device);
392        let batched_input =
393            Tensor::<TestBackend, 3>::random([8, 10, 64], Distribution::Default, &device);
394
395        let hidden_state = gru.forward(batched_input, None);
396
397        assert_eq!(hidden_state.shape().dims, [8, 10, 1024]);
398    }
399
400    #[test]
401    fn display() {
402        let config = GruConfig::new(2, 8, true);
403
404        let layer = config.init::<TestBackend>(&Default::default());
405
406        assert_eq!(
407            alloc::format!("{layer}"),
408            "Gru {d_input: 2, d_hidden: 8, bias: true, reset_after: true, params: 288}"
409        );
410    }
411}
burn_nn/modules/rnn/gru.rs

burn_nn/modules/rnn/
gru.rs
