trainers package

Submodules

trainers.classifier_trainer module

trainers.trainer module

class trainers.trainer.Trainer(diffusion_model, folder=None, *, dataset=None, train_batch_size=16, gradient_accumulate_every=1, augment_horizontal_flip=True, train_lr=0.0001, train_num_steps=100000, break_every_steps=None, ema_update_every=10, ema_decay=0.995, adam_betas=(0.9, 0.99), save_and_sample_every=1000, num_samples=25, results_folder='./results', amp=False, mixed_precision_type='fp16', split_batches=True, convert_image_to=None, calculate_fid=True, inception_block_idx=2048, max_grad_norm=1.0, num_fid_samples=50000, save_best_and_latest_only=False, tracker=None, tracker_kwargs=None)

Bases: Trainer

Parameters:

• diffusion_model (Module)

• folder (str)

• dataset (Dataset)

• train_batch_size (int)

• gradient_accumulate_every (int)

• augment_horizontal_flip (int)

• train_lr (float)

• train_num_steps (int)

• break_every_steps (int | None)

• ema_update_every (int)

• ema_decay (float)

• adam_betas (tuple[float, float])

• save_and_sample_every (int)

• num_samples (int)

• results_folder (str)

• amp (bool)

• mixed_precision_type (str)

• split_batches (bool)

• convert_image_to (Literal['L', 'RGB', 'RGBA'] | None)

• calculate_fid (bool)

• inception_block_idx (int)

• max_grad_norm (float)

• num_fid_samples (int)

• save_best_and_latest_only (bool)

• tracker (str | None)

• tracker_kwargs (dict | None)

save(milestone, keep_last_models=10)

Save the current state of the trainer.

Parameters:

• milestone (str) – The milestone identifier for the saved model.

• keep_last_models (int | None, optional) – The number of last models to keep. Defaults to 10.

keep_last_models(num_models=10)

Keep only the last num_models models in the results folder. Function will keep the latest model also.

Parameters:

num_models (int) – The number of models to keep.

Return type:

None

load(milestone_path)

Load the model, optimizer, and other training state from a milestone file.

Parameters:

milestone_path (str) – The path to the milestone file.

train()

Train the model.

This method performs the training loop for the model. It iterates over the specified number of training steps, accumulates gradients, updates the model parameters, and saves intermediate results if necessary.

Returns:

None

Module contents

This module provides the Trainer class for training models.

class trainers.Trainer(diffusion_model, folder=None, *, dataset=None, train_batch_size=16, gradient_accumulate_every=1, augment_horizontal_flip=True, train_lr=0.0001, train_num_steps=100000, break_every_steps=None, ema_update_every=10, ema_decay=0.995, adam_betas=(0.9, 0.99), save_and_sample_every=1000, num_samples=25, results_folder='./results', amp=False, mixed_precision_type='fp16', split_batches=True, convert_image_to=None, calculate_fid=True, inception_block_idx=2048, max_grad_norm=1.0, num_fid_samples=50000, save_best_and_latest_only=False, tracker=None, tracker_kwargs=None)

Bases: Trainer

Parameters:

• diffusion_model (Module)

• folder (str)

• dataset (Dataset)

• train_batch_size (int)

• gradient_accumulate_every (int)

• augment_horizontal_flip (int)

• train_lr (float)

• train_num_steps (int)

• break_every_steps (int | None)

• ema_update_every (int)

• ema_decay (float)

• adam_betas (tuple[float, float])

• save_and_sample_every (int)

• num_samples (int)

• results_folder (str)

• amp (bool)

• mixed_precision_type (str)

• split_batches (bool)

• convert_image_to (Literal['L', 'RGB', 'RGBA'] | None)

• calculate_fid (bool)

• inception_block_idx (int)

• max_grad_norm (float)

• num_fid_samples (int)

• save_best_and_latest_only (bool)

• tracker (str | None)

• tracker_kwargs (dict | None)

save(milestone, keep_last_models=10)

Save the current state of the trainer.

Parameters:

• milestone (str) – The milestone identifier for the saved model.

• keep_last_models (int | None, optional) – The number of last models to keep. Defaults to 10.

keep_last_models(num_models=10)

Keep only the last num_models models in the results folder. Function will keep the latest model also.

Parameters:

num_models (int) – The number of models to keep.

Return type:

None

load(milestone_path)

Load the model, optimizer, and other training state from a milestone file.

Parameters:

milestone_path (str) – The path to the milestone file.

train()

Train the model.

This method performs the training loop for the model. It iterates over the specified number of training steps, accumulates gradients, updates the model parameters, and saves intermediate results if necessary.

Returns:

None