pub struct Trainer { /* private fields */ }Expand description
Trainer for TinyRecursiveModel
Implementations§
Source§impl Trainer
impl Trainer
Sourcepub fn new(
model_config: TRMConfig,
training_config: TrainingConfig,
device: Device,
) -> Result<Self>
pub fn new( model_config: TRMConfig, training_config: TrainingConfig, device: Device, ) -> Result<Self>
Create new trainer
Sourcepub fn train_step(
&mut self,
input_ids: &Tensor,
target_ids: &Tensor,
) -> Result<f32>
pub fn train_step( &mut self, input_ids: &Tensor, target_ids: &Tensor, ) -> Result<f32>
Training step
Sourcepub fn save_checkpoint<P: AsRef<Path>>(
&self,
path: P,
loss: Option<f64>,
) -> Result<()>
pub fn save_checkpoint<P: AsRef<Path>>( &self, path: P, loss: Option<f64>, ) -> Result<()>
Save checkpoint
Sourcepub fn train_epoch(
&mut self,
dataloader: &mut impl BatchDataLoader,
) -> Result<f32>
pub fn train_epoch( &mut self, dataloader: &mut impl BatchDataLoader, ) -> Result<f32>
Train for one epoch
Sourcepub fn train(&mut self, dataloader: &mut impl BatchDataLoader) -> Result<()>
pub fn train(&mut self, dataloader: &mut impl BatchDataLoader) -> Result<()>
Full training loop
Auto Trait Implementations§
impl Freeze for Trainer
impl !RefUnwindSafe for Trainer
impl Send for Trainer
impl Sync for Trainer
impl Unpin for Trainer
impl !UnwindSafe for Trainer
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more
Source§impl<T> Instrument for T
impl<T> Instrument for T
Source§fn instrument(self, span: Span) -> Instrumented<Self>
fn instrument(self, span: Span) -> Instrumented<Self>
Source§fn in_current_span(self) -> Instrumented<Self>
fn in_current_span(self) -> Instrumented<Self>
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more