burn_train/learner/
train_val.rs

1use crate::{ItemLazy, renderer::MetricsRenderer};
2use burn_core::{module::Module, tensor::Gradients};
3use burn_optim::{
4    GradientsParams, ModuleOptimizer, MultiGradientsParams,
5    lr_scheduler::module_lr_scheduler::ModuleLearningRate,
6};
7
8/// A training output.
9pub struct TrainOutput<TO> {
10    /// The gradients.
11    pub grads: GradientsParams,
12
13    /// The item.
14    pub item: TO,
15}
16
17impl<TO> TrainOutput<TO> {
18    /// Creates a new training output.
19    ///
20    /// # Arguments
21    ///
22    /// * `module` - The module.
23    /// * `grads` - The gradients.
24    /// * `item` - The item.
25    ///
26    /// # Returns
27    ///
28    /// A new training output.
29    pub fn new<M: Module>(module: &M, grads: Gradients, item: TO) -> Self {
30        let grads = GradientsParams::from_grads(grads, module);
31        Self { grads, item }
32    }
33}
34
35/// Trait to be implemented for models to be able to be trained.
36///
37/// The [step](TrainStep::step) method needs to be manually implemented for all structs.
38///
39/// The [optimize](TrainStep::optimize) method can be overridden if you want to control how the
40/// optimizer is used to update the model. This can be useful if you want to call custom mutable
41/// functions on your model (e.g., clipping the weights) before or after the optimizer is used.
42///
43/// # Notes
44///
45/// To be used with the [Learner](crate::Learner) struct, the struct which implements this trait must
46/// also implement the [Module] trait, which is done automatically with its derive.
47pub trait TrainStep {
48    /// Type of input for a step of the training stage.
49    type Input: Send + 'static;
50    /// Type of output for a step of the training stage.
51    type Output: ItemLazy + 'static;
52    /// Runs a step for training, which executes the forward and backward passes.
53    ///
54    /// # Arguments
55    ///
56    /// * `item` - The input for the model.
57    ///
58    /// # Returns
59    ///
60    /// The output containing the model output and the gradients.
61    fn step(&self, item: Self::Input) -> TrainOutput<Self::Output>;
62    /// Optimize the current module with the provided gradients and learning rate.
63    ///
64    /// # Arguments
65    ///
66    /// * `optim`: Optimizer used for learning.
67    /// * `lr`: The learning rate used for this step.
68    /// * `grads`: The gradients of each parameter in the current model.
69    ///
70    /// # Returns
71    ///
72    /// The updated model.
73    fn optimize(
74        self,
75        optim: &mut ModuleOptimizer,
76        lr_module: ModuleLearningRate,
77        grads: GradientsParams,
78    ) -> Self
79    where
80        Self: Module + Sized,
81    {
82        optim.step(lr_module, self, grads)
83    }
84    /// Optimize the current module with the provided gradients and learning rate.
85    ///
86    /// # Arguments
87    ///
88    /// * `optim`: Optimizer used for learning.
89    /// * `lr`: The learning rate used for this step.
90    /// * `grads`: Multiple gradients associated to each parameter in the current model.
91    ///
92    /// # Returns
93    ///
94    /// The updated model.
95    fn optimize_multi(
96        self,
97        optim: &mut ModuleOptimizer,
98        lr_module: ModuleLearningRate,
99        grads: MultiGradientsParams,
100    ) -> Self
101    where
102        Self: Module + Sized,
103    {
104        optim.step_multi(lr_module, self, grads)
105    }
106}
107
108/// Trait to be implemented for validating models.
109pub trait InferenceStep {
110    /// Type of input for an inference step.
111    type Input: Send + 'static;
112    /// Type of output for an inference step.
113    type Output: ItemLazy + 'static;
114    /// Runs a validation step.
115    ///
116    /// # Arguments
117    ///
118    /// * `item` - The item to validate on.
119    ///
120    /// # Returns
121    ///
122    /// The validation output.
123    fn step(&self, item: Self::Input) -> Self::Output;
124}
125
126/// The result of a training, containing the model along with the [renderer](MetricsRenderer).
127pub struct LearningResult<M> {
128    /// The model with the learned weights, converted to validation mode.
129    ///
130    /// Call [`Module::train`](burn_core::module::Module::train) before using it for follow-up
131    /// training. To continue on another device, use `model.train().fork(device)`.
132    pub model: M,
133    /// The renderer that can be used for follow up training and evaluation.
134    pub renderer: Box<dyn MetricsRenderer>,
135    /// The stop that ended training before its last epoch, if
136    /// [`Interrupter::stop`](crate::Interrupter::stop) was called and no error happened.
137    pub interrupted: Option<crate::Interruption>,
138    /// The error that ended training before its last epoch, if it hit one.
139    pub error: Option<std::sync::Arc<crate::TrainingError>>,
140}
burn_train/learner/train_val.rs

burn_train/learner/
train_val.rs
