Trait TrainStep 

pub trait TrainStep {
    type Input: Send + 'static;
    type Output: ItemLazy + 'static;

    // Required method
    fn step(&self, item: Self::Input) -> TrainOutput<Self::Output>;

    // Provided methods
    fn optimize(
        self,
        optim: &mut ModuleOptimizer,
        lr_module: ModuleLearningRate,
        grads: GradientsParams,
    ) -> Self
       where Self: Sized + Module { ... }
    fn optimize_multi(
        self,
        optim: &mut ModuleOptimizer,
        lr_module: ModuleLearningRate,
        grads: MultiGradientsParams,
    ) -> Self
       where Self: Sized + Module { ... }
}

Trait to be implemented for models to be able to be trained.

The step method needs to be manually implemented for all structs.

The optimize method can be overridden if you want to control how the optimizer is used to update the model. This can be useful if you want to call custom mutable functions on your model (e.g., clipping the weights) before or after the optimizer is used.

Notes

To be used with the Learner struct, the struct which implements this trait must also implement the Module trait, which is done automatically with its derive.

Required Associated Types

type Input: Send + 'static

Type of input for a step of the training stage.

type Output: ItemLazy + 'static

Type of output for a step of the training stage.

Required Methods

fn step(&self, item: Self::Input) -> TrainOutput<Self::Output>

Runs a step for training, which executes the forward and backward passes.

Arguments

• item - The input for the model.

Returns

The output containing the model output and the gradients.

Provided Methods

fn optimize( self, optim: &mut ModuleOptimizer, lr_module: ModuleLearningRate, grads: GradientsParams, ) -> Self

where Self: Sized + Module,

Optimize the current module with the provided gradients and learning rate.

Arguments

• optim: Optimizer used for learning.
• lr: The learning rate used for this step.
• grads: The gradients of each parameter in the current model.

Returns

The updated model.

fn optimize_multi( self, optim: &mut ModuleOptimizer, lr_module: ModuleLearningRate, grads: MultiGradientsParams, ) -> Self

where Self: Sized + Module,

Optimize the current module with the provided gradients and learning rate.

Arguments

• optim: Optimizer used for learning.
• lr: The learning rate used for this step.
• grads: Multiple gradients associated to each parameter in the current model.

Returns

The updated model.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors