Struct Learner 

pub struct Learner<M: LearnerModel> { /* private fields */ }

Learner struct encapsulating all components necessary to train a Neural Network model.

Implementations

impl<M: LearnerModel> Learner<M>

pub fn new( model: M, optim: ModuleOptimizer, lr_scheduler: impl Into<ModuleLrScheduler>, ) -> Self

Create a learner.

impl<M: LearnerModel> Learner<M>

pub fn fork(&mut self, device: &Device)

Fork the learner’s model to the given device.

pub fn model(&self) -> M

Returns the current model.

pub fn lr_current(&self) -> ModuleLearningRate

Returns the current learning rate.

pub fn lr_step(&mut self)

Executes a step of the learning rate scheduler.

pub fn train_step( &self, item: <M as TrainStep>::Input, ) -> TrainOutput<<M as TrainStep>::Output>

Runs a step of the model 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.

pub fn optimizer_step(&mut self, grads: GradientsParams)

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.

pub fn optimizer_step_multi(&mut self, grads: MultiGradientsParams)

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.

pub fn load_model(&mut self, record: ModuleRecord)

Load the module state from a record.

pub fn load_optim(&mut self, record: OptimizerRecord)

Load the state of the learner’s optimizer from a record.

No device is needed: the optimizer state is migrated to each parameter’s device on the next step (see ModuleOptimizer::load_record).

pub fn load_scheduler(&mut self, record: LrSchedulerRecord)

Load the state of the learner’s scheduler from a record.

impl<M: LearnerModel> Learner<M>

pub fn grad_sharded(&mut self)

Mark the model as sharded across multiple devices.

Trait Implementations

impl<M: LearnerModel> Clone for Learner<M>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.