Skip to main content

rust_lstm/
lib.rs

1#![allow(clippy::type_complexity)]
2
3//! # Rust LSTM Library
4//!
5//! A complete LSTM implementation with training capabilities, multiple optimizers,
6//! dropout regularization, and support for various architectures including peephole
7//! connections and bidirectional processing.
8//!
9//! ## Core Components
10//!
11//! - **LSTM Cells**: Standard and peephole LSTM implementations with full backpropagation
12//! - **Bidirectional LSTM**: Process sequences in both directions with flexible output combination
13//! - **Networks**: Multi-layer LSTM networks for sequence modeling
14//! - **Training**: Complete training system with BPTT, gradient clipping, and validation
15//! - **Optimizers**: SGD, Adam, and RMSprop optimizers with adaptive learning rates
16//! - **Loss Functions**: MSE, MAE, and Cross-Entropy with numerically stable implementations
17//! - **Dropout**: Input, recurrent, output dropout and zoneout regularization
18//!
19//! ## Quick Start
20//!
21//! ```rust
22//! use rust_lstm::models::lstm_network::LSTMNetwork;
23//! use rust_lstm::training::create_basic_trainer;
24//!
25//! // Create a 2-layer LSTM with 10 input features and 20 hidden units
26//! let mut network = LSTMNetwork::new(10, 20, 2)
27//!     .with_input_dropout(0.2, true)     // Variational input dropout
28//!     .with_recurrent_dropout(0.3, true) // Variational recurrent dropout
29//!     .with_output_dropout(0.1);         // Standard output dropout
30//!
31//! let mut trainer = create_basic_trainer(network, 0.001);
32//!
33//! // Train on your data
34//! // trainer.train(&train_data, Some(&validation_data));
35//! ```
36
37pub mod layers;
38pub mod loss;
39pub mod models;
40pub mod optimizers;
41pub mod persistence;
42pub mod schedulers;
43pub mod text;
44pub mod training;
45/// Main library module.
46pub mod utils;
47
48// Re-export commonly used items
49pub use layers::bilstm_network::{BiLSTMNetwork, BiLSTMNetworkCache, CombineMode};
50pub use layers::dropout::{Dropout, Zoneout};
51pub use layers::gru_cell::{GRUCell, GRUCellCache, GRUCellGradients};
52pub use layers::linear::{LinearGradients, LinearLayer};
53pub use layers::lstm_cell::{LSTMCell, LSTMCellCache, LSTMCellGradients};
54pub use layers::peephole_lstm_cell::{
55    PeepholeLSTMCell, PeepholeLSTMCellCache, PeepholeLSTMCellGradients,
56};
57pub use loss::{CrossEntropyLoss, LossFunction, MAELoss, MSELoss};
58pub use models::gru_network::{
59    GRUNetwork, GRUNetworkCache, LayerDropoutConfig as GRULayerDropoutConfig,
60};
61pub use models::lstm_network::{LSTMNetwork, LSTMNetworkCache, LSTMState, LayerDropoutConfig};
62pub use optimizers::{Adam, RMSprop, ScheduledOptimizer, SGD};
63pub use persistence::{ModelMetadata, ModelPersistence, PersistenceError, PersistentModel};
64pub use schedulers::{
65    AnnealStrategy, ConstantLR, CosineAnnealingLR, CosineAnnealingWarmRestarts, CyclicalLR,
66    CyclicalMode, ExponentialLR, LRScheduleVisualizer, LearningRateScheduler, LinearLR,
67    MultiStepLR, OneCycleLR, PolynomialLR, ReduceLROnPlateau, ScaleMode, StepLR, WarmupScheduler,
68};
69pub use text::{
70    argmax, sample_nucleus, sample_top_k, sample_with_temperature, softmax, CharacterEmbedding,
71    EmbeddingGradients, TextVocabulary,
72};
73pub use training::{
74    create_adam_batch_trainer, create_basic_batch_trainer, create_basic_trainer,
75    create_cosine_annealing_trainer, create_one_cycle_trainer, create_step_lr_trainer,
76    EarlyStopper, EarlyStoppingConfig, EarlyStoppingMetric, LSTMBatchTrainer, LSTMTrainer,
77    ScheduledLSTMTrainer, TrainingConfig, TrainingMetrics,
78};
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use ndarray::arr2;
84
85    #[test]
86    fn test_library_integration() {
87        let mut network = models::lstm_network::LSTMNetwork::new(2, 3, 1);
88        let input = arr2(&[[1.0], [0.5]]);
89        let state = network.zero_state(1);
90
91        let (hy, next) = network.forward(&input, &state);
92        let cy = &next.c[0];
93
94        assert_eq!(hy.shape(), &[3, 1]);
95        assert_eq!(cy.shape(), &[3, 1]);
96    }
97
98    #[test]
99    fn test_library_with_dropout() {
100        let mut network = models::lstm_network::LSTMNetwork::new(2, 3, 1)
101            .with_input_dropout(0.2, false)
102            .with_recurrent_dropout(0.3, true)
103            .with_output_dropout(0.1);
104
105        let input = arr2(&[[1.0], [0.5]]);
106        let state = network.zero_state(1);
107
108        // Test training mode
109        network.train();
110        let (hy_train, next_train) = network.forward(&input, &state);
111        let cy_train = &next_train.c[0];
112
113        // Test evaluation mode
114        network.eval();
115        let (hy_eval, next_eval) = network.forward(&input, &state);
116        let cy_eval = &next_eval.c[0];
117
118        assert_eq!(hy_train.shape(), &[3, 1]);
119        assert_eq!(cy_train.shape(), &[3, 1]);
120        assert_eq!(hy_eval.shape(), &[3, 1]);
121        assert_eq!(cy_eval.shape(), &[3, 1]);
122    }
123}