rust_lstm/
training.rs

1use crate::layers::lstm_cell::LSTMCellGradients;
2use crate::loss::{LossFunction, MSELoss};
3use crate::models::lstm_network::{pad_batch, LSTMNetwork};
4use crate::optimizers::{Optimizer, ScheduledOptimizer, SGD};
5use crate::persistence::SerializableLSTMNetwork;
6use crate::schedulers::LearningRateScheduler;
7use ndarray::Array2;
8use std::time::Instant;
9
10/// Loss summed over the time steps of one sequence and its gradients obtained by
11/// backpropagation through time.
12pub fn sequence_gradients<L: LossFunction>(
13    network: &mut LSTMNetwork,
14    loss_function: &L,
15    inputs: &[Array2<f64>],
16    targets: &[Array2<f64>],
17) -> (f64, Vec<LSTMCellGradients>) {
18    let (outputs, caches) = network.forward_sequence_with_cache(inputs);
19    let mut total_loss = 0.0;
20    let d_outputs: Vec<Array2<f64>> = outputs
21        .iter()
22        .zip(targets)
23        .map(|((output, _), target)| {
24            total_loss += loss_function.compute_loss(output, target);
25            loss_function.compute_gradient(output, target)
26        })
27        .collect();
28    let (gradients, _) = network.backward_sequence(&d_outputs, &caches);
29    (total_loss, gradients)
30}
31
32/// Rescales all gradients together so that their joint L2 norm is at most `max_norm`.
33pub fn clip_global_norm(gradients: &mut [LSTMCellGradients], max_norm: f64) {
34    let norm = gradients
35        .iter()
36        .map(LSTMCellGradients::squared_norm)
37        .sum::<f64>()
38        .sqrt();
39    if norm > max_norm {
40        for g in gradients.iter_mut() {
41            g.scale(max_norm / norm);
42        }
43    }
44}
45
46/// Configuration for training hyperparameters
47pub struct TrainingConfig {
48    pub epochs: usize,
49    pub print_every: usize,
50    pub clip_gradient: Option<f64>,
51    pub log_lr_changes: bool,
52    pub early_stopping: Option<EarlyStoppingConfig>,
53}
54
55/// Configuration for early stopping
56#[derive(Debug, Clone)]
57pub struct EarlyStoppingConfig {
58    /// Number of epochs with no improvement after which training will be stopped
59    pub patience: usize,
60    /// Minimum change in the monitored quantity to qualify as an improvement
61    pub min_delta: f64,
62    /// Whether to restore the best weights when early stopping triggers
63    pub restore_best_weights: bool,
64    /// Metric to monitor for early stopping ('val_loss' or 'train_loss')
65    pub monitor: EarlyStoppingMetric,
66}
67
68/// Metric to monitor for early stopping
69#[derive(Debug, Clone, PartialEq)]
70pub enum EarlyStoppingMetric {
71    ValidationLoss,
72    TrainLoss,
73}
74
75impl Default for EarlyStoppingConfig {
76    fn default() -> Self {
77        EarlyStoppingConfig {
78            patience: 10,
79            min_delta: 1e-4,
80            restore_best_weights: true,
81            monitor: EarlyStoppingMetric::ValidationLoss,
82        }
83    }
84}
85
86impl Default for TrainingConfig {
87    fn default() -> Self {
88        TrainingConfig {
89            epochs: 100,
90            print_every: 10,
91            clip_gradient: Some(5.0),
92            log_lr_changes: true,
93            early_stopping: None,
94        }
95    }
96}
97
98/// Training metrics tracked during training
99#[derive(Debug, Clone)]
100pub struct TrainingMetrics {
101    pub epoch: usize,
102    pub train_loss: f64,
103    pub validation_loss: Option<f64>,
104    pub time_elapsed: f64,
105    pub learning_rate: f64,
106}
107
108/// Early stopping state tracker
109#[derive(Debug, Clone)]
110pub struct EarlyStopper {
111    config: EarlyStoppingConfig,
112    best_score: f64,
113    wait_count: usize,
114    stopped_epoch: Option<usize>,
115    best_weights: Option<SerializableLSTMNetwork>, // Serialized network weights
116}
117
118impl EarlyStopper {
119    pub fn new(config: EarlyStoppingConfig) -> Self {
120        EarlyStopper {
121            config,
122            best_score: f64::INFINITY,
123            wait_count: 0,
124            stopped_epoch: None,
125            best_weights: None,
126        }
127    }
128
129    /// Check if training should stop based on current metrics
130    /// Returns (should_stop, is_best_score)
131    pub fn should_stop(
132        &mut self,
133        current_metrics: &TrainingMetrics,
134        network: &LSTMNetwork,
135    ) -> (bool, bool) {
136        let current_score = match self.config.monitor {
137            EarlyStoppingMetric::ValidationLoss => {
138                match current_metrics.validation_loss {
139                    Some(val_loss) => val_loss,
140                    None => {
141                        // If validation loss is not available, fall back to train loss
142                        current_metrics.train_loss
143                    }
144                }
145            }
146            EarlyStoppingMetric::TrainLoss => current_metrics.train_loss,
147        };
148
149        let is_improvement = current_score < self.best_score - self.config.min_delta;
150
151        if is_improvement {
152            self.best_score = current_score;
153            self.wait_count = 0;
154
155            // Save best weights if restore_best_weights is enabled
156            if self.config.restore_best_weights {
157                self.best_weights = Some(network.into());
158            }
159
160            (false, true)
161        } else {
162            self.wait_count += 1;
163
164            if self.wait_count >= self.config.patience {
165                self.stopped_epoch = Some(current_metrics.epoch);
166                (true, false)
167            } else {
168                (false, false)
169            }
170        }
171    }
172
173    /// Get the epoch where training was stopped
174    pub fn stopped_epoch(&self) -> Option<usize> {
175        self.stopped_epoch
176    }
177
178    /// Get the best score achieved
179    pub fn best_score(&self) -> f64 {
180        self.best_score
181    }
182
183    /// Restore the best weights to the network if available
184    pub fn restore_best_weights(&self, network: &mut LSTMNetwork) -> Result<(), String> {
185        if let Some(ref weights) = self.best_weights {
186            *network = weights.clone().into();
187            Ok(())
188        } else {
189            Err("No best weights available to restore".to_string())
190        }
191    }
192}
193
194/// Main trainer for LSTM networks with configurable loss and optimizer
195pub struct LSTMTrainer<L: LossFunction, O: Optimizer> {
196    pub network: LSTMNetwork,
197    pub loss_function: L,
198    pub optimizer: O,
199    pub config: TrainingConfig,
200    pub metrics_history: Vec<TrainingMetrics>,
201    early_stopper: Option<EarlyStopper>,
202}
203
204impl<L: LossFunction, O: Optimizer> LSTMTrainer<L, O> {
205    pub fn new(network: LSTMNetwork, loss_function: L, optimizer: O) -> Self {
206        LSTMTrainer {
207            network,
208            loss_function,
209            optimizer,
210            config: TrainingConfig::default(),
211            metrics_history: Vec::new(),
212            early_stopper: None,
213        }
214    }
215
216    pub fn with_config(mut self, config: TrainingConfig) -> Self {
217        // Initialize early stopper if early stopping is configured
218        self.early_stopper = config
219            .early_stopping
220            .as_ref()
221            .map(|es_config| EarlyStopper::new(es_config.clone()));
222        self.config = config;
223        self
224    }
225
226    /// Train on a single sequence using backpropagation through time (BPTT)
227    pub fn train_sequence(&mut self, inputs: &[Array2<f64>], targets: &[Array2<f64>]) -> f64 {
228        if inputs.len() != targets.len() {
229            panic!("Inputs and targets must have the same length");
230        }
231
232        self.network.train();
233
234        let (total_loss, mut total_gradients) =
235            sequence_gradients(&mut self.network, &self.loss_function, inputs, targets);
236
237        if let Some(clip_value) = self.config.clip_gradient {
238            clip_global_norm(&mut total_gradients, clip_value);
239        }
240
241        self.network
242            .update_parameters(&total_gradients, &mut self.optimizer);
243
244        total_loss / inputs.len() as f64
245    }
246
247    /// Train for multiple epochs with optional validation
248    pub fn train(
249        &mut self,
250        train_data: &[(Vec<Array2<f64>>, Vec<Array2<f64>>)],
251        validation_data: Option<&[(Vec<Array2<f64>>, Vec<Array2<f64>>)]>,
252    ) {
253        println!("Starting training for {} epochs...", self.config.epochs);
254
255        for epoch in 0..self.config.epochs {
256            let start_time = Instant::now();
257            let mut epoch_loss = 0.0;
258
259            // Training phase
260            self.network.train();
261            for (inputs, targets) in train_data {
262                let loss = self.train_sequence(inputs, targets);
263                epoch_loss += loss;
264            }
265            epoch_loss /= train_data.len() as f64;
266
267            let validation_loss = if let Some(val_data) = validation_data {
268                self.network.eval();
269                Some(self.evaluate(val_data))
270            } else {
271                None
272            };
273
274            let time_elapsed = start_time.elapsed().as_secs_f64();
275
276            let current_lr = self.optimizer.get_learning_rate();
277            let metrics = TrainingMetrics {
278                epoch,
279                train_loss: epoch_loss,
280                validation_loss,
281                time_elapsed,
282                learning_rate: current_lr,
283            };
284
285            self.metrics_history.push(metrics.clone());
286
287            // Check early stopping
288            let mut should_stop = false;
289            let mut is_best = false;
290            if let Some(ref mut early_stopper) = self.early_stopper {
291                let (stop, best) = early_stopper.should_stop(&metrics, &self.network);
292                should_stop = stop;
293                is_best = best;
294            }
295
296            if epoch % self.config.print_every == 0 {
297                let best_indicator = if is_best { " *" } else { "" };
298                if let Some(val_loss) = validation_loss {
299                    println!(
300                        "Epoch {}: Train Loss: {:.6}, Val Loss: {:.6}, LR: {:.2e}, Time: {:.2}s{}",
301                        epoch, epoch_loss, val_loss, current_lr, time_elapsed, best_indicator
302                    );
303                } else {
304                    println!(
305                        "Epoch {}: Train Loss: {:.6}, LR: {:.2e}, Time: {:.2}s{}",
306                        epoch, epoch_loss, current_lr, time_elapsed, best_indicator
307                    );
308                }
309            }
310
311            if should_stop {
312                let stopped_epoch = self
313                    .early_stopper
314                    .as_ref()
315                    .unwrap()
316                    .stopped_epoch()
317                    .unwrap();
318                let best_score = self.early_stopper.as_ref().unwrap().best_score();
319                println!(
320                    "Early stopping triggered at epoch {} (best score: {:.6})",
321                    stopped_epoch, best_score
322                );
323
324                // Restore best weights if configured
325                if let Some(ref early_stopper) = self.early_stopper {
326                    if let Err(e) = early_stopper.restore_best_weights(&mut self.network) {
327                        println!("Warning: Could not restore best weights: {}", e);
328                    } else {
329                        println!(
330                            "Restored best weights from epoch with score {:.6}",
331                            best_score
332                        );
333                    }
334                }
335                break;
336            }
337        }
338
339        println!("Training completed!");
340    }
341
342    /// Evaluate model performance on validation data
343    pub fn evaluate(&mut self, data: &[(Vec<Array2<f64>>, Vec<Array2<f64>>)]) -> f64 {
344        self.network.eval();
345
346        let mut total_loss = 0.0;
347        let mut total_samples = 0;
348
349        for (inputs, targets) in data {
350            if inputs.len() != targets.len() {
351                continue;
352            }
353
354            let (outputs, _) = self.network.forward_sequence_with_cache(inputs);
355
356            for ((output, _), target) in outputs.iter().zip(targets.iter()) {
357                let loss = self.loss_function.compute_loss(output, target);
358                total_loss += loss;
359                total_samples += 1;
360            }
361        }
362
363        if total_samples > 0 {
364            total_loss / total_samples as f64
365        } else {
366            0.0
367        }
368    }
369
370    /// Generate predictions for input sequences
371    pub fn predict(&mut self, inputs: &[Array2<f64>]) -> Vec<Array2<f64>> {
372        self.network.eval();
373
374        let (outputs, _) = self.network.forward_sequence_with_cache(inputs);
375        outputs.into_iter().map(|(output, _)| output).collect()
376    }
377
378    pub fn get_latest_metrics(&self) -> Option<&TrainingMetrics> {
379        self.metrics_history.last()
380    }
381
382    pub fn get_metrics_history(&self) -> &[TrainingMetrics] {
383        &self.metrics_history
384    }
385
386    /// Set network to training mode
387    pub fn set_training_mode(&mut self, training: bool) {
388        if training {
389            self.network.train();
390        } else {
391            self.network.eval();
392        }
393    }
394}
395
396/// Specialized trainer for scheduled optimizers that automatically steps the scheduler
397pub struct ScheduledLSTMTrainer<L: LossFunction, O: Optimizer, S: LearningRateScheduler> {
398    pub network: LSTMNetwork,
399    pub loss_function: L,
400    pub optimizer: ScheduledOptimizer<O, S>,
401    pub config: TrainingConfig,
402    pub metrics_history: Vec<TrainingMetrics>,
403    early_stopper: Option<EarlyStopper>,
404}
405
406impl<L: LossFunction, O: Optimizer, S: LearningRateScheduler> ScheduledLSTMTrainer<L, O, S> {
407    pub fn new(
408        network: LSTMNetwork,
409        loss_function: L,
410        optimizer: ScheduledOptimizer<O, S>,
411    ) -> Self {
412        ScheduledLSTMTrainer {
413            network,
414            loss_function,
415            optimizer,
416            config: TrainingConfig::default(),
417            metrics_history: Vec::new(),
418            early_stopper: None,
419        }
420    }
421
422    pub fn with_config(mut self, config: TrainingConfig) -> Self {
423        // Initialize early stopper if early stopping is configured
424        self.early_stopper = config
425            .early_stopping
426            .as_ref()
427            .map(|es_config| EarlyStopper::new(es_config.clone()));
428        self.config = config;
429        self
430    }
431
432    /// Train on a single sequence using backpropagation through time (BPTT)
433    pub fn train_sequence(&mut self, inputs: &[Array2<f64>], targets: &[Array2<f64>]) -> f64 {
434        if inputs.len() != targets.len() {
435            panic!("Inputs and targets must have the same length");
436        }
437
438        self.network.train();
439
440        let (total_loss, mut total_gradients) =
441            sequence_gradients(&mut self.network, &self.loss_function, inputs, targets);
442
443        if let Some(clip_value) = self.config.clip_gradient {
444            clip_global_norm(&mut total_gradients, clip_value);
445        }
446
447        self.network
448            .update_parameters(&total_gradients, &mut self.optimizer);
449
450        total_loss / inputs.len() as f64
451    }
452
453    /// Train for multiple epochs with automatic scheduler stepping
454    pub fn train(
455        &mut self,
456        train_data: &[(Vec<Array2<f64>>, Vec<Array2<f64>>)],
457        validation_data: Option<&[(Vec<Array2<f64>>, Vec<Array2<f64>>)]>,
458    ) {
459        println!(
460            "Starting training for {} epochs with {} scheduler...",
461            self.config.epochs,
462            self.optimizer.scheduler_name()
463        );
464
465        for epoch in 0..self.config.epochs {
466            let start_time = Instant::now();
467            let mut epoch_loss = 0.0;
468
469            // Training phase
470            self.network.train();
471            for (inputs, targets) in train_data {
472                let loss = self.train_sequence(inputs, targets);
473                epoch_loss += loss;
474            }
475            epoch_loss /= train_data.len() as f64;
476
477            let validation_loss = if let Some(val_data) = validation_data {
478                self.network.eval();
479                Some(self.evaluate(val_data))
480            } else {
481                None
482            };
483
484            // Step the scheduler at the end of each epoch
485            let prev_lr = self.optimizer.get_learning_rate();
486            if let Some(val_loss) = validation_loss {
487                self.optimizer.step_with_val_loss(val_loss);
488            } else {
489                self.optimizer.step();
490            }
491            let new_lr = self.optimizer.get_learning_rate();
492
493            // Log learning rate changes if enabled
494            if self.config.log_lr_changes && (new_lr - prev_lr).abs() > 1e-10 {
495                println!(
496                    "Learning rate changed from {:.2e} to {:.2e}",
497                    prev_lr, new_lr
498                );
499            }
500
501            let time_elapsed = start_time.elapsed().as_secs_f64();
502
503            let metrics = TrainingMetrics {
504                epoch,
505                train_loss: epoch_loss,
506                validation_loss,
507                time_elapsed,
508                learning_rate: new_lr,
509            };
510
511            self.metrics_history.push(metrics.clone());
512
513            // Check early stopping
514            let mut should_stop = false;
515            let mut is_best = false;
516            if let Some(ref mut early_stopper) = self.early_stopper {
517                let (stop, best) = early_stopper.should_stop(&metrics, &self.network);
518                should_stop = stop;
519                is_best = best;
520            }
521
522            if epoch % self.config.print_every == 0 {
523                let best_indicator = if is_best { " *" } else { "" };
524                if let Some(val_loss) = validation_loss {
525                    println!(
526                        "Epoch {}: Train Loss: {:.6}, Val Loss: {:.6}, LR: {:.2e}, Time: {:.2}s{}",
527                        epoch, epoch_loss, val_loss, new_lr, time_elapsed, best_indicator
528                    );
529                } else {
530                    println!(
531                        "Epoch {}: Train Loss: {:.6}, LR: {:.2e}, Time: {:.2}s{}",
532                        epoch, epoch_loss, new_lr, time_elapsed, best_indicator
533                    );
534                }
535            }
536
537            if should_stop {
538                let stopped_epoch = self
539                    .early_stopper
540                    .as_ref()
541                    .unwrap()
542                    .stopped_epoch()
543                    .unwrap();
544                let best_score = self.early_stopper.as_ref().unwrap().best_score();
545                println!(
546                    "Early stopping triggered at epoch {} (best score: {:.6})",
547                    stopped_epoch, best_score
548                );
549
550                // Restore best weights if configured
551                if let Some(ref early_stopper) = self.early_stopper {
552                    if let Err(e) = early_stopper.restore_best_weights(&mut self.network) {
553                        println!("Warning: Could not restore best weights: {}", e);
554                    } else {
555                        println!(
556                            "Restored best weights from epoch with score {:.6}",
557                            best_score
558                        );
559                    }
560                }
561                break;
562            }
563        }
564
565        println!("Training completed!");
566    }
567
568    /// Evaluate model performance on validation data
569    pub fn evaluate(&mut self, data: &[(Vec<Array2<f64>>, Vec<Array2<f64>>)]) -> f64 {
570        self.network.eval();
571
572        let mut total_loss = 0.0;
573        let mut total_samples = 0;
574
575        for (inputs, targets) in data {
576            if inputs.len() != targets.len() {
577                continue;
578            }
579
580            let (outputs, _) = self.network.forward_sequence_with_cache(inputs);
581
582            for ((output, _), target) in outputs.iter().zip(targets.iter()) {
583                let loss = self.loss_function.compute_loss(output, target);
584                total_loss += loss;
585                total_samples += 1;
586            }
587        }
588
589        if total_samples > 0 {
590            total_loss / total_samples as f64
591        } else {
592            0.0
593        }
594    }
595
596    /// Generate predictions for input sequences
597    pub fn predict(&mut self, inputs: &[Array2<f64>]) -> Vec<Array2<f64>> {
598        self.network.eval();
599
600        let (outputs, _) = self.network.forward_sequence_with_cache(inputs);
601        outputs.into_iter().map(|(output, _)| output).collect()
602    }
603
604    pub fn get_latest_metrics(&self) -> Option<&TrainingMetrics> {
605        self.metrics_history.last()
606    }
607
608    pub fn get_metrics_history(&self) -> &[TrainingMetrics] {
609        &self.metrics_history
610    }
611
612    /// Set network to training mode
613    pub fn set_training_mode(&mut self, training: bool) {
614        if training {
615            self.network.train();
616        } else {
617            self.network.eval();
618        }
619    }
620
621    /// Get the current learning rate
622    pub fn get_current_lr(&self) -> f64 {
623        self.optimizer.get_current_lr()
624    }
625
626    /// Get the current epoch from the scheduler
627    pub fn get_current_epoch(&self) -> usize {
628        self.optimizer.get_current_epoch()
629    }
630
631    /// Reset the optimizer and scheduler
632    pub fn reset_optimizer(&mut self) {
633        self.optimizer.reset();
634    }
635}
636
637/// Batch trainer for LSTM networks with configurable loss and optimizer
638/// Processes multiple sequences simultaneously for improved performance
639pub struct LSTMBatchTrainer<L: LossFunction, O: Optimizer> {
640    pub network: LSTMNetwork,
641    pub loss_function: L,
642    pub optimizer: O,
643    pub config: TrainingConfig,
644    pub metrics_history: Vec<TrainingMetrics>,
645    early_stopper: Option<EarlyStopper>,
646}
647
648impl<L: LossFunction, O: Optimizer> LSTMBatchTrainer<L, O> {
649    pub fn new(network: LSTMNetwork, loss_function: L, optimizer: O) -> Self {
650        LSTMBatchTrainer {
651            network,
652            loss_function,
653            optimizer,
654            config: TrainingConfig::default(),
655            metrics_history: Vec::new(),
656            early_stopper: None,
657        }
658    }
659
660    pub fn with_config(mut self, config: TrainingConfig) -> Self {
661        // Initialize early stopper if early stopping is configured
662        self.early_stopper = config
663            .early_stopping
664            .as_ref()
665            .map(|es_config| EarlyStopper::new(es_config.clone()));
666        self.config = config;
667        self
668    }
669
670    /// Train on a batch of sequences using batch processing
671    ///
672    /// # Arguments
673    /// * `batch_inputs` - Vector of input sequences, each sequence is `Vec<Array2<f64>>`
674    /// * `batch_targets` - Vector of target sequences, each sequence is `Vec<Array2<f64>>`
675    ///
676    /// # Returns
677    /// * Average loss across the batch
678    pub fn train_batch(
679        &mut self,
680        batch_inputs: &[Vec<Array2<f64>>],
681        batch_targets: &[Vec<Array2<f64>>],
682    ) -> f64 {
683        assert_eq!(
684            batch_inputs.len(),
685            batch_targets.len(),
686            "Batch inputs and targets must have same length"
687        );
688
689        if batch_inputs.is_empty() {
690            return 0.0;
691        }
692
693        self.network.train();
694
695        let batch_size = batch_inputs.len();
696        let hidden_size = self.network.hidden_size;
697        let padded_inputs = pad_batch(batch_inputs, self.network.input_size);
698        let (outputs, caches) = self.network.forward_sequence_with_cache(&padded_inputs);
699
700        let mut total_loss = 0.0;
701        let mut valid_steps = 0;
702        let mut d_outputs = Vec::with_capacity(outputs.len());
703        for (t, (prediction, _)) in outputs.iter().enumerate() {
704            let active: Vec<usize> = (0..batch_size)
705                .filter(|&b| t < batch_inputs[b].len() && t < batch_targets[b].len())
706                .collect();
707            let mut d_output = Array2::zeros((hidden_size, batch_size));
708            if !active.is_empty() {
709                let mut preds = Array2::zeros((hidden_size, active.len()));
710                let mut targs = Array2::zeros((hidden_size, active.len()));
711                for (i, &b) in active.iter().enumerate() {
712                    preds.column_mut(i).assign(&prediction.column(b));
713                    targs.column_mut(i).assign(&batch_targets[b][t].column(0));
714                }
715                total_loss += self.loss_function.compute_batch_loss(&preds, &targs);
716                valid_steps += 1;
717                let grad = self.loss_function.compute_batch_gradient(&preds, &targs);
718                for (i, &b) in active.iter().enumerate() {
719                    d_output.column_mut(b).assign(&grad.column(i));
720                }
721            }
722            d_outputs.push(d_output);
723        }
724
725        let (mut total_gradients, _) = self.network.backward_sequence(&d_outputs, &caches);
726
727        if let Some(clip_value) = self.config.clip_gradient {
728            clip_global_norm(&mut total_gradients, clip_value);
729        }
730
731        // Update parameters
732        self.network
733            .update_parameters(&total_gradients, &mut self.optimizer);
734
735        if valid_steps > 0 {
736            total_loss / valid_steps as f64
737        } else {
738            0.0
739        }
740    }
741
742    /// Train for multiple epochs with batch processing
743    ///
744    /// # Arguments
745    /// * `train_data` - Vector of (input_sequences, target_sequences) tuples for training
746    /// * `validation_data` - Optional validation data
747    /// * `batch_size` - Number of sequences to process in each batch
748    pub fn train(
749        &mut self,
750        train_data: &[(Vec<Array2<f64>>, Vec<Array2<f64>>)],
751        validation_data: Option<&[(Vec<Array2<f64>>, Vec<Array2<f64>>)]>,
752        batch_size: usize,
753    ) {
754        println!(
755            "Starting batch training for {} epochs with batch size {}...",
756            self.config.epochs, batch_size
757        );
758
759        for epoch in 0..self.config.epochs {
760            let start_time = Instant::now();
761            let mut epoch_loss = 0.0;
762            let mut num_batches = 0;
763
764            // Create batches
765            for batch_start in (0..train_data.len()).step_by(batch_size) {
766                let batch_end = (batch_start + batch_size).min(train_data.len());
767                let batch = &train_data[batch_start..batch_end];
768
769                let batch_inputs: Vec<_> = batch.iter().map(|(inputs, _)| inputs.clone()).collect();
770                let batch_targets: Vec<_> =
771                    batch.iter().map(|(_, targets)| targets.clone()).collect();
772
773                let batch_loss = self.train_batch(&batch_inputs, &batch_targets);
774                epoch_loss += batch_loss;
775                num_batches += 1;
776            }
777
778            epoch_loss /= num_batches as f64;
779
780            // Validation
781            let validation_loss = if let Some(val_data) = validation_data {
782                self.network.eval();
783                Some(self.evaluate_batch(val_data, batch_size))
784            } else {
785                None
786            };
787
788            let time_elapsed = start_time.elapsed().as_secs_f64();
789            let current_lr = self.optimizer.get_learning_rate();
790
791            let metrics = TrainingMetrics {
792                epoch,
793                train_loss: epoch_loss,
794                validation_loss,
795                time_elapsed,
796                learning_rate: current_lr,
797            };
798
799            self.metrics_history.push(metrics.clone());
800
801            // Check early stopping
802            let mut should_stop = false;
803            let mut is_best = false;
804            if let Some(ref mut early_stopper) = self.early_stopper {
805                let (stop, best) = early_stopper.should_stop(&metrics, &self.network);
806                should_stop = stop;
807                is_best = best;
808            }
809
810            if epoch % self.config.print_every == 0 {
811                let best_indicator = if is_best { " *" } else { "" };
812                if let Some(val_loss) = validation_loss {
813                    println!("Epoch {}: Train Loss: {:.6}, Val Loss: {:.6}, LR: {:.2e}, Time: {:.2}s, Batches: {}{}", 
814                             epoch, epoch_loss, val_loss, current_lr, time_elapsed, num_batches, best_indicator);
815                } else {
816                    println!(
817                        "Epoch {}: Train Loss: {:.6}, LR: {:.2e}, Time: {:.2}s, Batches: {}{}",
818                        epoch, epoch_loss, current_lr, time_elapsed, num_batches, best_indicator
819                    );
820                }
821            }
822
823            if should_stop {
824                let stopped_epoch = self
825                    .early_stopper
826                    .as_ref()
827                    .unwrap()
828                    .stopped_epoch()
829                    .unwrap();
830                let best_score = self.early_stopper.as_ref().unwrap().best_score();
831                println!(
832                    "Early stopping triggered at epoch {} (best score: {:.6})",
833                    stopped_epoch, best_score
834                );
835
836                // Restore best weights if configured
837                if let Some(ref early_stopper) = self.early_stopper {
838                    if let Err(e) = early_stopper.restore_best_weights(&mut self.network) {
839                        println!("Warning: Could not restore best weights: {}", e);
840                    } else {
841                        println!(
842                            "Restored best weights from epoch with score {:.6}",
843                            best_score
844                        );
845                    }
846                }
847                break;
848            }
849        }
850
851        println!("Batch training completed!");
852    }
853
854    /// Evaluate model performance using batch processing
855    pub fn evaluate_batch(
856        &mut self,
857        data: &[(Vec<Array2<f64>>, Vec<Array2<f64>>)],
858        batch_size: usize,
859    ) -> f64 {
860        self.network.eval();
861
862        let mut total_loss = 0.0;
863        let mut num_batches = 0;
864
865        for batch_start in (0..data.len()).step_by(batch_size) {
866            let batch_end = (batch_start + batch_size).min(data.len());
867            let batch = &data[batch_start..batch_end];
868
869            let batch_inputs: Vec<_> = batch.iter().map(|(inputs, _)| inputs.clone()).collect();
870            let batch_targets: Vec<_> = batch.iter().map(|(_, targets)| targets.clone()).collect();
871
872            // Process batch and compute loss (simplified evaluation)
873            let batch_outputs = self.network.forward_batch_sequences(&batch_inputs);
874
875            let mut batch_loss = 0.0;
876            let mut valid_samples = 0;
877
878            for (outputs, targets) in batch_outputs.iter().zip(batch_targets.iter()) {
879                for ((output, _), target) in outputs.iter().zip(targets.iter()) {
880                    let loss = self.loss_function.compute_loss(output, target);
881                    batch_loss += loss;
882                    valid_samples += 1;
883                }
884            }
885
886            if valid_samples > 0 {
887                total_loss += batch_loss / valid_samples as f64;
888                num_batches += 1;
889            }
890        }
891
892        if num_batches > 0 {
893            total_loss / num_batches as f64
894        } else {
895            0.0
896        }
897    }
898
899    /// Generate predictions using batch processing
900    pub fn predict_batch(&mut self, inputs: &[Vec<Array2<f64>>]) -> Vec<Vec<Array2<f64>>> {
901        self.network.eval();
902
903        let batch_outputs = self.network.forward_batch_sequences(inputs);
904        batch_outputs
905            .into_iter()
906            .map(|sequence_outputs| {
907                sequence_outputs
908                    .into_iter()
909                    .map(|(output, _)| output)
910                    .collect()
911            })
912            .collect()
913    }
914
915    pub fn get_latest_metrics(&self) -> Option<&TrainingMetrics> {
916        self.metrics_history.last()
917    }
918
919    pub fn get_metrics_history(&self) -> &[TrainingMetrics] {
920        &self.metrics_history
921    }
922
923    pub fn set_training_mode(&mut self, training: bool) {
924        if training {
925            self.network.train();
926        } else {
927            self.network.eval();
928        }
929    }
930}
931
932/// Create a basic trainer with SGD optimizer and MSE loss
933pub fn create_basic_trainer(network: LSTMNetwork, learning_rate: f64) -> LSTMTrainer<MSELoss, SGD> {
934    let loss_function = MSELoss;
935    let optimizer = SGD::new(learning_rate);
936    LSTMTrainer::new(network, loss_function, optimizer)
937}
938
939/// Create a scheduled trainer with SGD and StepLR scheduler
940pub fn create_step_lr_trainer(
941    network: LSTMNetwork,
942    learning_rate: f64,
943    step_size: usize,
944    gamma: f64,
945) -> ScheduledLSTMTrainer<MSELoss, SGD, crate::schedulers::StepLR> {
946    let loss_function = MSELoss;
947    let optimizer =
948        ScheduledOptimizer::step_lr(SGD::new(learning_rate), learning_rate, step_size, gamma);
949    ScheduledLSTMTrainer::new(network, loss_function, optimizer)
950}
951
952/// Create a scheduled trainer with Adam and OneCycleLR scheduler  
953pub fn create_one_cycle_trainer(
954    network: LSTMNetwork,
955    max_lr: f64,
956    total_steps: usize,
957) -> ScheduledLSTMTrainer<MSELoss, crate::optimizers::Adam, crate::schedulers::OneCycleLR> {
958    let loss_function = MSELoss;
959    let optimizer =
960        ScheduledOptimizer::one_cycle(crate::optimizers::Adam::new(max_lr), max_lr, total_steps);
961    ScheduledLSTMTrainer::new(network, loss_function, optimizer)
962}
963
964/// Create a scheduled trainer with Adam and CosineAnnealingLR scheduler
965pub fn create_cosine_annealing_trainer(
966    network: LSTMNetwork,
967    learning_rate: f64,
968    t_max: usize,
969    eta_min: f64,
970) -> ScheduledLSTMTrainer<MSELoss, crate::optimizers::Adam, crate::schedulers::CosineAnnealingLR> {
971    let loss_function = MSELoss;
972    let optimizer = crate::optimizers::Adam::new(learning_rate);
973    let scheduler = crate::schedulers::CosineAnnealingLR::new(t_max, eta_min);
974    let scheduled_optimizer =
975        crate::optimizers::ScheduledOptimizer::new(optimizer, scheduler, learning_rate);
976
977    ScheduledLSTMTrainer::new(network, loss_function, scheduled_optimizer)
978}
979
980/// Create a basic batch trainer with SGD optimizer and MSE loss
981pub fn create_basic_batch_trainer(
982    network: LSTMNetwork,
983    learning_rate: f64,
984) -> LSTMBatchTrainer<MSELoss, SGD> {
985    let loss_function = MSELoss;
986    let optimizer = SGD::new(learning_rate);
987    LSTMBatchTrainer::new(network, loss_function, optimizer)
988}
989
990/// Create a batch trainer with Adam optimizer and MSE loss
991pub fn create_adam_batch_trainer(
992    network: LSTMNetwork,
993    learning_rate: f64,
994) -> LSTMBatchTrainer<MSELoss, crate::optimizers::Adam> {
995    let loss_function = MSELoss;
996    let optimizer = crate::optimizers::Adam::new(learning_rate);
997    LSTMBatchTrainer::new(network, loss_function, optimizer)
998}
999
1000#[cfg(test)]
1001mod tests {
1002    use super::*;
1003    use ndarray::arr2;
1004
1005    #[test]
1006    fn test_trainer_creation() {
1007        let network = LSTMNetwork::new(2, 3, 1);
1008        let trainer = create_basic_trainer(network, 0.01);
1009
1010        assert_eq!(trainer.network.input_size, 2);
1011        assert_eq!(trainer.network.hidden_size, 3);
1012        assert_eq!(trainer.network.num_layers, 1);
1013    }
1014
1015    #[test]
1016    fn test_sequence_training() {
1017        let network = LSTMNetwork::new(2, 3, 1);
1018        let mut trainer = create_basic_trainer(network, 0.01);
1019
1020        let inputs = vec![arr2(&[[1.0], [0.0]]), arr2(&[[0.0], [1.0]])];
1021        let targets = vec![arr2(&[[1.0], [0.0], [0.0]]), arr2(&[[0.0], [1.0], [0.0]])];
1022
1023        let loss = trainer.train_sequence(&inputs, &targets);
1024        assert!(loss >= 0.0);
1025    }
1026}
rust_lstm/training.rs

rust_lstm/
training.rs
