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

rust_lstm/models/
lstm_network.rs
