Scheduling within models

[tomhosking]

Some models make use of some sort of scheduling or annealing internally (e.g. KL warmup or temperature annealing) based on the current step index - what's the correct way to implement this within pythae?

[clementchadebec]

Hi @tomhosking,

Thank you for opening this issue. Yes indeed some models may need several warmup steps during training. To allow such a behavior, the current training epoch is automatically passed to the forward method of the model in the trainers as shown below

benchmark_VAE/src/pythae/trainers/base_trainer/base_trainer.py

Lines 473 to 475 in 0f5c0cc

	model_output = self.model(
	inputs, epoch=epoch, dataset_size=len(self.train_loader.dataset)
	)

The current epoch can be retrieved in the forward method as follows

epoch = kwargs.pop("epoch")

and then passed to the loss_function. Nonetheless, this warmup setting is only supported for some models for now and you may require to make the changes in the model implementation yourself. If you do so, you will need to first clone the repo and install it in editable mode

$ git clone https://github.com/clementchadebec/benchmark_VAE.git
$ cd benchmark_VAE
$ pip install -e .

As an example, you can check the implementation of the disentangled_beta_vae that uses a specific KL scheduling procedure.

I will think of a way to include this for other models as well.

Best,

Clément