1#![allow(clippy::type_complexity)]
2
3pub mod layers;
38pub mod loss;
39pub mod models;
40pub mod optimizers;
41pub mod persistence;
42pub mod schedulers;
43pub mod text;
44pub mod training;
45pub mod utils;
47
48pub 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 network.train();
110 let (hy_train, next_train) = network.forward(&input, &state);
111 let cy_train = &next_train.c[0];
112
113 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}