MultistageTrainer

Trainer for multistage training. It is used for models that have a reset_optimizer_epochs attribute: JNF, TELBO.

class multivae.trainers.MultistageTrainerConfig(output_dir=None, per_device_train_batch_size=64, per_device_eval_batch_size=64, num_epochs=100, train_dataloader_num_workers=0, eval_dataloader_num_workers=0, optimizer_cls='Adam', optimizer_params=None, scheduler_cls=None, scheduler_params=None, learning_rate=0.0001, steps_saving=None, steps_predict=None, keep_best_on_train=False, seed=8, no_cuda=False, world_size=-1, local_rank=-1, rank=-1, dist_backend='nccl', master_addr='localhost', master_port='12345', drop_last=False, gradient_clipping_max_norm=None)

Configuration for a specific trainer that handles the training of the joint VAE models.

Parameters:

• output_dir (str) – The directory where model checkpoints, configs and final model will be stored. Default: None.

• per_device_train_batch_size (int) – The number of training samples per batch and per device. Default 64

• per_device_eval_batch_size (int) – The number of evaluation samples per batch and per device. Default 64

• num_epochs (int) – The maximal number of epochs for training. Default: 100

• train_dataloader_num_workers (int) – Number of subprocesses to use for train data loading. 0 means that the data will be loaded in the main process. Default: 0

• eval_dataloader_num_workers (int) – Number of subprocesses to use for evaluation data loading. 0 means that the data will be loaded in the main process. Default: 0

• optimizer_cls (str) – The name of the torch.optim.Optimizer used for training. Default: Adam.

• optimizer_params (dict) – A dict containing the parameters to use for the torch.optim.Optimizer. If None, uses the default parameters. Default: None.

• scheduler_cls (str) – The name of the torch.optim.lr_scheduler used for training. If None, no scheduler is used. Default None.

• scheduler_params (dict) – A dict containing the parameters to use for the torch.optim.le_scheduler. If None, uses the default parameters. Default: None.

• learning_rate (int) – The learning rate applied to the Optimizer. Default: 1e-4

• steps_saving (int) – A model checkpoint will be saved every steps_saving epoch. Default: None

• steps_predict (int) – A prediction using the best model will be run every steps_predict epoch. Default: None

• keep_best_on_train (bool) – Whether to keep the best model on the train set. Default: False

• seed (int) – The random seed for reproducibility

• no_cuda (bool) – Disable cuda training. Default: False

• world_size (int) – The total number of process to run. Default: -1

• local_rank (int) – The rank of the node for distributed training. Default: -1

• rank (int) – The rank of the process for distributed training. Default: -1

• dist_backend (str) – The distributed backend to use. Default: ‘nccl’

• master_addr (str) – The master address for distributed training. Default: ‘localhost’

• master_port (str) – The master port for distributed training. Default: ‘12345’

class multivae.trainers.MultistageTrainer(model, train_dataset, eval_dataset=None, training_config=None, callbacks=None, checkpoint=None)

A specific trainer that handles the training of the joint VAE models.

Parameters:

• model (BaseMultiVAE) – A instance of BaseMultiVAE to train

• train_dataset (MultimodalBaseDataset) – The training dataset of type MultimodalBaseDataset

• eval_dataset (MultimodalBaseDataset) – The evaluation dataset of type MultimodalBaseDataset

• training_config (BaseTrainerConfig) – The training arguments summarizing the main parameters used for training. If None, a basic training instance of BaseTrainerConfig is used. Default: None.

• callbacks (List[TrainingCallback]) – A list of callbacks to use during training.

prepare_train_step(epoch, best_train_loss, best_eval_loss)

Function to operate changes between train_steps such as resetting the optimizer and the best losses values.