Struct burn::train::TrainEpoch

pub struct TrainEpoch<TI> { /* private fields */ }

A training epoch.

Implementations

impl<TI> TrainEpoch<TI>

pub fn new( dataloader: Arc<dyn DataLoader<TI>>, epoch: usize, epoch_total: usize, grad_accumulation: Option<usize>, ) -> TrainEpoch<TI>

Constructs a new TrainEpoch.

impl<TI> TrainEpoch<TI>

pub fn run<LC, TO>( &self, model: <LC as LearnerComponents>::Model, optim: <LC as LearnerComponents>::Optimizer, scheduler: &mut <LC as LearnerComponents>::LrScheduler, processor: &mut <LC as LearnerComponents>::EventProcessor, interrupter: &TrainingInterrupter, ) -> (<LC as LearnerComponents>::Model, <LC as LearnerComponents>::Optimizer)

where LC: LearnerComponents, <LC as LearnerComponents>::EventProcessor: EventProcessor<ItemTrain = TO>, <LC as LearnerComponents>::Model: TrainStep<TI, TO>,

Runs the training epoch.

Arguments

• model - The model to train.
• optim - The optimizer to use.
• scheduler - The learning rate scheduler to use.
• processor - The event processor to use.

Returns

The trained model and the optimizer.

impl<TI> TrainEpoch<TI>

pub fn run_multi_device<LC, TO>( &self, model: <LC as LearnerComponents>::Model, optim: <LC as LearnerComponents>::Optimizer, lr_scheduler: &mut <LC as LearnerComponents>::LrScheduler, processor: &mut <LC as LearnerComponents>::EventProcessor, devices: Vec<<<LC as LearnerComponents>::Backend as Backend>::Device>, interrupter: &TrainingInterrupter, ) -> (<LC as LearnerComponents>::Model, <LC as LearnerComponents>::Optimizer)

where LC: LearnerComponents, <LC as LearnerComponents>::EventProcessor: EventProcessor<ItemTrain = TO>, <LC as LearnerComponents>::Model: TrainStep<TI, TO>, TO: Send + 'static, TI: Send + 'static,

Runs the training epoch on multiple devices.

Arguments

• model - The model to train.
• optim - The optimizer to use.
• lr_scheduler - The learning rate scheduler to use.
• processor - The event processor to use.
• devices - The devices to use.

Returns

The trained model and the optimizer.