rust_lstm/models/
lstm_network.rs

1use crate::layers::lstm_cell::{masked, LSTMCell, LSTMCellCache, LSTMCellGradients};
2use crate::optimizers::Optimizer;
3use ndarray::Array2;
4
5/// Hidden and cell states of every layer, each of shape (hidden_size, batch).
6#[derive(Clone, Debug, PartialEq)]
7pub struct LSTMState {
8    pub h: Vec<Array2<f64>>,
9    pub c: Vec<Array2<f64>>,
10}
11
12/// Values saved by one network time step.
13#[derive(Clone)]
14pub struct LSTMNetworkCache {
15    pub cell_caches: Vec<LSTMCellCache>,
16    /// Inverted-dropout masks applied between layer `i` and layer `i + 1`.
17    pub output_dropout_masks: Vec<Option<Array2<f64>>>,
18}
19
20/// Multi-layer LSTM network for sequence modeling with dropout support
21///
22/// Stacks multiple LSTM cells where the output of layer i becomes
23/// the input to layer i+1. Supports both inference and training with
24/// configurable dropout regularization.
25#[derive(Clone)]
26pub struct LSTMNetwork {
27    cells: Vec<LSTMCell>,
28    pub input_size: usize,
29    pub hidden_size: usize,
30    pub num_layers: usize,
31    pub is_training: bool,
32}
33
34impl LSTMNetwork {
35    /// Creates a new multi-layer LSTM network
36    ///
37    /// First layer accepts `input_size` dimensions, subsequent layers
38    /// accept `hidden_size` dimensions from the previous layer.
39    pub fn new(input_size: usize, hidden_size: usize, num_layers: usize) -> Self {
40        let mut cells = Vec::new();
41
42        for i in 0..num_layers {
43            let layer_input_size = if i == 0 { input_size } else { hidden_size };
44            cells.push(LSTMCell::new(layer_input_size, hidden_size));
45        }
46
47        LSTMNetwork {
48            cells,
49            input_size,
50            hidden_size,
51            num_layers,
52            is_training: true,
53        }
54    }
55
56    pub fn with_input_dropout(mut self, dropout_rate: f64, variational: bool) -> Self {
57        for cell in &mut self.cells {
58            *cell = cell.clone().with_input_dropout(dropout_rate, variational);
59        }
60        self
61    }
62
63    pub fn with_recurrent_dropout(mut self, dropout_rate: f64, variational: bool) -> Self {
64        for cell in &mut self.cells {
65            *cell = cell
66                .clone()
67                .with_recurrent_dropout(dropout_rate, variational);
68        }
69        self
70    }
71
72    pub fn with_output_dropout(mut self, dropout_rate: f64) -> Self {
73        for (i, cell) in self.cells.iter_mut().enumerate() {
74            if i < self.num_layers - 1 {
75                *cell = cell.clone().with_output_dropout(dropout_rate);
76            }
77        }
78        self
79    }
80
81    pub fn with_zoneout(mut self, cell_zoneout_rate: f64, hidden_zoneout_rate: f64) -> Self {
82        for cell in &mut self.cells {
83            *cell = cell
84                .clone()
85                .with_zoneout(cell_zoneout_rate, hidden_zoneout_rate);
86        }
87        self
88    }
89
90    /// Dropout on the candidate update of every layer (Semeniuta et al., 2016).
91    pub fn with_cell_update_dropout(mut self, dropout_rate: f64, variational: bool) -> Self {
92        for cell in &mut self.cells {
93            *cell = cell
94                .clone()
95                .with_cell_update_dropout(dropout_rate, variational);
96        }
97        self
98    }
99
100    pub fn with_layer_dropout(mut self, layer_configs: Vec<LayerDropoutConfig>) -> Self {
101        for (i, config) in layer_configs.into_iter().enumerate() {
102            if i < self.cells.len() {
103                let mut cell = self.cells[i].clone();
104
105                if let Some((rate, variational)) = config.input_dropout {
106                    cell = cell.with_input_dropout(rate, variational);
107                }
108                if let Some((rate, variational)) = config.recurrent_dropout {
109                    cell = cell.with_recurrent_dropout(rate, variational);
110                }
111                if let Some(rate) = config.output_dropout {
112                    cell = cell.with_output_dropout(rate);
113                }
114                if let Some((cell_rate, hidden_rate)) = config.zoneout {
115                    cell = cell.with_zoneout(cell_rate, hidden_rate);
116                }
117                if let Some((rate, variational)) = config.cell_update_dropout {
118                    cell = cell.with_cell_update_dropout(rate, variational);
119                }
120
121                self.cells[i] = cell;
122            }
123        }
124        self
125    }
126
127    pub fn train(&mut self) {
128        self.is_training = true;
129        for cell in &mut self.cells {
130            cell.train();
131        }
132    }
133
134    pub fn eval(&mut self) {
135        self.is_training = false;
136        for cell in &mut self.cells {
137            cell.eval();
138        }
139    }
140
141    /// Creates a network from existing cells (used for deserialization)
142    pub fn from_cells(
143        cells: Vec<LSTMCell>,
144        input_size: usize,
145        hidden_size: usize,
146        num_layers: usize,
147    ) -> Self {
148        LSTMNetwork {
149            cells,
150            input_size,
151            hidden_size,
152            num_layers,
153            is_training: true,
154        }
155    }
156
157    /// Get reference to the cells (used for serialization)
158    pub fn get_cells(&self) -> &[LSTMCell] {
159        &self.cells
160    }
161
162    /// Get mutable reference to the cells (for training mode changes)
163    pub fn get_cells_mut(&mut self) -> &mut [LSTMCell] {
164        &mut self.cells
165    }
166
167    /// All-zero state for `batch_size` sequences.
168    pub fn zero_state(&self, batch_size: usize) -> LSTMState {
169        let zeros = || {
170            (0..self.num_layers)
171                .map(|_| Array2::zeros((self.hidden_size, batch_size)))
172                .collect()
173        };
174        LSTMState {
175            h: zeros(),
176            c: zeros(),
177        }
178    }
179
180    /// Clears variational dropout masks so that the next sequence samples new ones.
181    pub fn reset_dropout_masks(&mut self) {
182        for cell in &mut self.cells {
183            cell.reset_dropout_masks();
184        }
185    }
186
187    /// One time step. `input` has shape (input_size, batch). Returns the output of
188    /// the top layer and the new state of every layer.
189    pub fn forward(&mut self, input: &Array2<f64>, state: &LSTMState) -> (Array2<f64>, LSTMState) {
190        let (output, state, _) = self.forward_with_cache(input, state);
191        (output, state)
192    }
193
194    /// One time step that also returns the values needed by `backward_sequence`.
195    pub fn forward_with_cache(
196        &mut self,
197        input: &Array2<f64>,
198        state: &LSTMState,
199    ) -> (Array2<f64>, LSTMState, LSTMNetworkCache) {
200        assert_eq!(
201            state.h.len(),
202            self.num_layers,
203            "state must hold one entry per layer"
204        );
205        let mut layer_input = input.clone();
206        let mut next = LSTMState {
207            h: Vec::with_capacity(self.num_layers),
208            c: Vec::with_capacity(self.num_layers),
209        };
210        let mut cell_caches = Vec::with_capacity(self.num_layers);
211        let mut output_dropout_masks = Vec::with_capacity(self.num_layers);
212
213        let last = self.num_layers - 1;
214        for (l, cell) in self.cells.iter_mut().enumerate() {
215            let (hy, cy, cache) = cell.forward_with_cache(&layer_input, &state.h[l], &state.c[l]);
216            let mask = if l < last {
217                cell.output_dropout_mask(hy.raw_dim())
218            } else {
219                None
220            };
221            layer_input = masked(&hy, &mask);
222            next.h.push(hy);
223            next.c.push(cy);
224            cell_caches.push(cache);
225            output_dropout_masks.push(mask);
226        }
227
228        let cache = LSTMNetworkCache {
229            cell_caches,
230            output_dropout_masks,
231        };
232        (layer_input, next, cache)
233    }
234
235    /// Runs a whole sequence from a zero state. Each element of `sequence` has shape
236    /// (input_size, batch); returns the top-layer output at every step.
237    pub fn forward_sequence(&mut self, sequence: &[Array2<f64>]) -> Vec<Array2<f64>> {
238        self.forward_sequence_with_cache(sequence)
239            .0
240            .into_iter()
241            .map(|(h, _)| h)
242            .collect()
243    }
244
245    /// Runs a whole sequence from a zero state and keeps the caches for
246    /// `backward_sequence`. Returns the top-layer hidden and cell state at every step.
247    pub fn forward_sequence_with_cache(
248        &mut self,
249        sequence: &[Array2<f64>],
250    ) -> (Vec<(Array2<f64>, Array2<f64>)>, Vec<LSTMNetworkCache>) {
251        let batch_size = sequence.first().map_or(1, |x| x.ncols());
252        let mut state = self.zero_state(batch_size);
253        self.reset_dropout_masks();
254
255        let mut outputs = Vec::with_capacity(sequence.len());
256        let mut caches = Vec::with_capacity(sequence.len());
257        for input in sequence {
258            let (output, next, cache) = self.forward_with_cache(input, &state);
259            outputs.push((output, next.c[self.num_layers - 1].clone()));
260            caches.push(cache);
261            state = next;
262        }
263        (outputs, caches)
264    }
265
266    /// Backpropagation through time over a sequence processed by
267    /// `forward_sequence_with_cache`.
268    ///
269    /// `d_outputs[t]` is the gradient of the loss with respect to the top-layer
270    /// output at step `t`. Returns the parameter gradients of every layer, summed
271    /// over time and batch, and the gradient with respect to every input.
272    pub fn backward_sequence(
273        &self,
274        d_outputs: &[Array2<f64>],
275        caches: &[LSTMNetworkCache],
276    ) -> (Vec<LSTMCellGradients>, Vec<Array2<f64>>) {
277        assert_eq!(
278            d_outputs.len(),
279            caches.len(),
280            "one output gradient per time step"
281        );
282        let mut gradients = self.zero_gradients();
283        let mut d_inputs = vec![Array2::zeros((0, 0)); caches.len()];
284        let batch_size = d_outputs.first().map_or(1, |d| d.ncols());
285        let mut dh_next = self.zero_state(batch_size);
286
287        let last = self.num_layers - 1;
288        for t in (0..caches.len()).rev() {
289            let mut d_above = d_outputs[t].clone();
290            for l in (0..self.num_layers).rev() {
291                let d_out = if l < last {
292                    masked(&d_above, &caches[t].output_dropout_masks[l])
293                } else {
294                    d_above
295                };
296                let dh = &d_out + &dh_next.h[l];
297                let (g, dx, dhx, dcx) =
298                    self.cells[l].backward(&dh, &dh_next.c[l], &caches[t].cell_caches[l]);
299                gradients[l].accumulate(&g);
300                dh_next.h[l] = dhx;
301                dh_next.c[l] = dcx;
302                d_above = dx;
303            }
304            d_inputs[t] = d_above;
305        }
306        (gradients, d_inputs)
307    }
308
309    /// Update parameters for all layers using computed gradients
310    pub fn update_parameters<O: Optimizer>(
311        &mut self,
312        gradients: &[LSTMCellGradients],
313        optimizer: &mut O,
314    ) {
315        for (i, (cell, cell_gradients)) in self.cells.iter_mut().zip(gradients.iter()).enumerate() {
316            let prefix = format!("layer_{}", i);
317            cell.update_parameters(cell_gradients, optimizer, &prefix);
318        }
319    }
320
321    /// Initialize zero gradients for all layers
322    pub fn zero_gradients(&self) -> Vec<LSTMCellGradients> {
323        self.cells
324            .iter()
325            .map(|cell| cell.zero_gradients())
326            .collect()
327    }
328
329    /// Processes sequences of different lengths as one batch.
330    ///
331    /// Every element of every sequence has shape (input_size, 1). Shorter sequences
332    /// are padded with zeros; their outputs after the last real step are discarded.
333    pub fn forward_batch_sequences(
334        &mut self,
335        batch_sequences: &[Vec<Array2<f64>>],
336    ) -> Vec<Vec<(Array2<f64>, Array2<f64>)>> {
337        let padded = pad_batch(batch_sequences, self.input_size);
338        let (outputs, _) = self.forward_sequence_with_cache(&padded);
339        batch_sequences
340            .iter()
341            .enumerate()
342            .map(|(b, seq)| {
343                outputs
344                    .iter()
345                    .take(seq.len())
346                    .map(|(h, c)| (column(h, b), column(c, b)))
347                    .collect()
348            })
349            .collect()
350    }
351}
352
353/// Stacks sequences of column vectors into zero-padded (features, batch) matrices.
354pub fn pad_batch(batch_sequences: &[Vec<Array2<f64>>], features: usize) -> Vec<Array2<f64>> {
355    let max_len = batch_sequences.iter().map(Vec::len).max().unwrap_or(0);
356    (0..max_len)
357        .map(|t| {
358            let mut step = Array2::zeros((features, batch_sequences.len()));
359            for (b, seq) in batch_sequences.iter().enumerate() {
360                if let Some(x) = seq.get(t) {
361                    step.column_mut(b).assign(&x.column(0));
362                }
363            }
364            step
365        })
366        .collect()
367}
368
369fn column(m: &Array2<f64>, b: usize) -> Array2<f64> {
370    m.column(b).to_owned().insert_axis(ndarray::Axis(1))
371}
372
373/// Configuration for layer-specific dropout settings
374#[derive(Clone, Debug)]
375pub struct LayerDropoutConfig {
376    pub input_dropout: Option<(f64, bool)>, // (rate, variational)
377    pub recurrent_dropout: Option<(f64, bool)>, // (rate, variational)
378    pub output_dropout: Option<f64>,        // rate
379    pub zoneout: Option<(f64, f64)>,        // (cell_rate, hidden_rate)
380    pub cell_update_dropout: Option<(f64, bool)>, // (rate, variational)
381}
382
383impl Default for LayerDropoutConfig {
384    fn default() -> Self {
385        Self::new()
386    }
387}
388
389impl LayerDropoutConfig {
390    pub fn new() -> Self {
391        LayerDropoutConfig {
392            input_dropout: None,
393            recurrent_dropout: None,
394            output_dropout: None,
395            zoneout: None,
396            cell_update_dropout: None,
397        }
398    }
399
400    pub fn with_input_dropout(mut self, rate: f64, variational: bool) -> Self {
401        self.input_dropout = Some((rate, variational));
402        self
403    }
404
405    pub fn with_recurrent_dropout(mut self, rate: f64, variational: bool) -> Self {
406        self.recurrent_dropout = Some((rate, variational));
407        self
408    }
409
410    pub fn with_output_dropout(mut self, rate: f64) -> Self {
411        self.output_dropout = Some(rate);
412        self
413    }
414
415    pub fn with_zoneout(mut self, cell_rate: f64, hidden_rate: f64) -> Self {
416        self.zoneout = Some((cell_rate, hidden_rate));
417        self
418    }
419
420    pub fn with_cell_update_dropout(mut self, rate: f64, variational: bool) -> Self {
421        self.cell_update_dropout = Some((rate, variational));
422        self
423    }
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429    use ndarray::arr2;
430
431    #[test]
432    fn test_lstm_network_forward() {
433        let input_size = 3;
434        let hidden_size = 2;
435        let num_layers = 2;
436        let mut network = LSTMNetwork::new(input_size, hidden_size, num_layers);
437
438        let input = arr2(&[[0.5], [0.1], [-0.3]]);
439        let state = network.zero_state(1);
440
441        let (hy, next) = network.forward(&input, &state);
442        let cy = next.c[num_layers - 1].clone();
443
444        assert_eq!(hy.shape(), &[hidden_size, 1]);
445        assert_eq!(cy.shape(), &[hidden_size, 1]);
446    }
447
448    #[test]
449    fn test_lstm_network_with_dropout() {
450        let input_size = 3;
451        let hidden_size = 2;
452        let num_layers = 2;
453        let mut network = LSTMNetwork::new(input_size, hidden_size, num_layers)
454            .with_input_dropout(0.2, true) // Variational input dropout
455            .with_recurrent_dropout(0.3, true) // Variational recurrent dropout
456            .with_output_dropout(0.1)
457            .with_zoneout(0.1, 0.1);
458
459        let input = arr2(&[[0.5], [0.1], [-0.3]]);
460        let state = network.zero_state(1);
461
462        // Test training mode
463        network.train();
464        let (hy_train, next) = network.forward(&input, &state);
465        let cy_train = next.c[num_layers - 1].clone();
466
467        // Test evaluation mode
468        network.eval();
469        let (hy_eval, next) = network.forward(&input, &state);
470        let cy_eval = next.c[num_layers - 1].clone();
471
472        assert_eq!(hy_train.shape(), &[hidden_size, 1]);
473        assert_eq!(cy_train.shape(), &[hidden_size, 1]);
474        assert_eq!(hy_eval.shape(), &[hidden_size, 1]);
475        assert_eq!(cy_eval.shape(), &[hidden_size, 1]);
476    }
477
478    #[test]
479    fn test_layer_specific_dropout() {
480        let input_size = 3;
481        let hidden_size = 2;
482        let num_layers = 2;
483
484        let layer_configs = vec![
485            LayerDropoutConfig::new()
486                .with_input_dropout(0.2, true)
487                .with_recurrent_dropout(0.3, true),
488            LayerDropoutConfig::new()
489                .with_output_dropout(0.1)
490                .with_zoneout(0.1, 0.1),
491        ];
492
493        let mut network =
494            LSTMNetwork::new(input_size, hidden_size, num_layers).with_layer_dropout(layer_configs);
495
496        let input = arr2(&[[0.5], [0.1], [-0.3]]);
497        let state = network.zero_state(1);
498
499        let (hy, next) = network.forward(&input, &state);
500        let cy = next.c[num_layers - 1].clone();
501
502        assert_eq!(hy.shape(), &[hidden_size, 1]);
503        assert_eq!(cy.shape(), &[hidden_size, 1]);
504    }
505}
rust_lstm/models/lstm_network.rs

rust_lstm/models/
lstm_network.rs
