[Question] Resume from checkpoint clarification

[lucalazzaroni]

Describe your question

I'm currently working on a transfer learning project for ASR and I followed the related tutorial. Since I'm not able to perform an entire training in a single session, I need to use checkpoints and resume from them.
The step I'm doing are the following:

1. Save the model after training:

quartznet.save_to("./out/my_model.nemo")

2. Restore the model when resuming:

quartznet = nemo_asr.models.EncDecCTCModel.restore_from("./out/my_model.nemo")

3. Restore the pytorch-lightning trainer:

logger = TensorBoardLogger(
    save_dir=os.getcwd(),
    version=3,
    name='lightning_logs'
)
trainer = pl.Trainer(
    gpus=1, max_epochs=30, precision=16, amp_level='O1', checkpoint_callback=True,         
    resume_from_checkpoint='lightning_logs/version_3/checkpoints/epoch=18-step=68893.ckpt', logger=logger)

4. Assign parameters and trainer to the model again:

# Set learning rate
params['model']['optim']['lr'] = 0.001
# Set NovoGrad optimizer betas
params['model']['optim']['betas'] = [0.95, 0.25]
# Set CosineAnnealing learning rate policy's warmup ratio
params['model']['optim']['sched']['warmup_ratio'] = 0.12
# Set training and validation labels
params['model']['train_ds']['labels'] = italian_vocabulary
params['model']['validation_ds']['labels'] = italian_vocabulary
# Set batch size
params['model']['train_ds']['batch_size'] = 16
params['model']['validation_ds']['batch_size'] = 16

# Assign trainer to the model
quartznet.set_trainer(trainer)

# Point to the data we'll use for fine-tuning as the training set
quartznet.setup_training_data(train_data_config=params['model']['train_ds'])

# Point to the new validation data for fine-tuning
quartznet.setup_validation_data(val_data_config=params['model']['validation_ds'])

# Add changes to quartznet model
quartznet.setup_optimization(optim_config=DictConfig(params['model']['optim']))

5. Resume training:

trainer.fit(quartznet)

I would like to know if this is the correct way to resume from checkpoint with the old trainer and model trained before, or if I'm missing some steps (or making redundant ones).
Before I tried to restore only the model and declare a new trainer (trainer = pl.Trainer(gpus=1, max_epochs=20, precision=16, amp_level='O1', checkpoint_callback=True)), then train for the same number of epochs of the first training session, but in the end I noticed the resulting WER was roughly the same, so I assumed the problem was that I hadn't restored my old trainer too.

Environment overview (please complete the following information)

• Environment location: Bare-metal
• Method of NeMo install: pip install git+https://github.com/NVIDIA/NeMo.git@main#egg=nemo_toolkit[all]

Environment details

• OS version: Ubuntu 18.04.5 LTS
• PyTorch version: pytorch lightning 1.1.5
• Python version: 3.6.9

Additional context

GPU model: NVIDIA Quadro RTX 4000
CUDA version: 10.1

[titu1994]

This is a great question ! Surprisingly, the QuartzNet config does not show two very nice parameters of experiment manager that are used precisely for resuming training over multiple runs - https://github.com/NVIDIA/NeMo/blob/main/examples/asr/conf/citrinet/citrinet_384.yaml#L381-L382

You can add them to the quartznet config manually, or use hydra to add them with the +exp_manager.resume_if_exists=true and +exp_manager.resume_ignore_no_checkpoint=true flags. Note, when doing multiple such runs, give the experiment directory a unique name with exp_manager.name=<some unique experiment name>.

After the first run, simply rerun the training script with the same config / overrides and same experiment manager flags - and it will load up the model + its checkpoint + the optimizer state and scheduler state automatically and continue training.

[titu1994]

A little explanation of what those two flags do

resume_if_exists = Resume the experiment if an earlier checkpoint exists. Docstring

resume_ignore_no_checkpoint = If you use the above flag alone for the first run, (when no checkpoint exists yet to resume from), experiment manager will throw an error (since it cant resume from a checkpoint which doesnt exist). This second flag tells exp_manager to ignore that the checkpoint doesnt exist for the first run - so the first run will "train from scratch". When it finishes, subsequent calls will see that a checkpoint exists so experiment manager will use it. Docstring

[lucalazzaroni]

Thanks for your detailed response!
I tried a training following the steps I listed before in my question and it seems not to improve: learning rate stays at 0, probably because Cosine Annealing took it to 0 at the end of the first training, but someway, if the old trainer is re-used increasing its max_epochs flag, the lr value is not increased anymore. So I don't think mine was the correct way.

Now I'm trying adding the parameters you suggested to the quartznet config. If I understood correctly, after the first run I can resume from checkpoint with quartznet = nemo_asr.models.EncDecCTCModel.restore_from("./out/my_model.nemo"), then re-assign experiment manager flags and overrides (i.e. quartznet.setup_training_data(train_data_config=params['model']['train_ds']) and so on, as in point 4 of my post) and, in the end, run trainer.fit(quartznet) without necessarily having to restore my old trainer as I did until now (point 3 of my post) but being able to create a new trainer such as trainer = pl.Trainer(gpus=1, max_epochs=20, precision=16, amp_level='O1'). Is it correct?

[lucalazzaroni]

Another question is, when I setup my trainer, should I give it a max_epochs flag equal to the total number of epochs I intend to play (i.e. 256) or a smaller number which is the number of epochs to be played during the current session (around 20)? Because I think this would influence the learning rate decay

[titu1994]

Please see the examples/asr/speech_to_text.py script in NeMo. You can replace the constructor EncDecCTCModel(...) With model = EncDecCTCModel.restore_from().

You will also need to use model.set_trainer(trainer), model.setup_training_data(cfg.model.train_ds), model.setup_validation_data(cfg.model.validation_ds) since restored models don't automatically setup their datasets and trainer.

After that, simply calling the script multiple times with the suggested experimental manager settings should be sufficient.

Note - after the first run, exp manager will create .ckpt Pytorch Lightning checkpoints - and will restore those for the next run. .nemo files are "final packaged" models - not intermediate training checkpoints. They don't contain the optimizer state etc, just the model parameters and the artifacts necessary to build the model.

Once your training finishes, exp manager will create a final . nemo file for you as well to show that training is over.

[titu1994]

Max epochs is the number of epochs it will train - and will affect the scheduler + learning rate.

[lucalazzaroni]

Thanks again for your answers. I followed the steps you suggested and now it seems to correctly resume from last checkpoint. The only thing I'm missing regards the number of epochs to be passed to the trainer. In my last experiment I set it to 256 (which is my target epochs value) and I manually stopped the trainer around the 20th epoch, then resumed after. Would it be different if I instead set the max_epochs value to 20 then increase it whenever I want to resume the training? I bet on setting it to 256 since you told me that max_epochs affects the scheduler + learning rate, but I'm not sure if manually stopping the training is the correct way or not.

[titu1994]

Once you begin a run, the history of what parameters you used at the beginning is preserved inside the .ckpt intermediate checkpoints. So you cannot increase number of epochs as such after the first run - it will complete it according to whatever was the original config only.

When restarting training, experiement manager will use whichever was the last checkpoint saved. If you stop training between two save steps, the subsrqu run will have to again train on the steps that were not saved prior to being stopped.

[lucalazzaroni]

Many thanks, you have clarified all my doubts. I think we can close the issue.

[kzelias]

You will also need to use model.set_trainer(trainer), model.setup_training_data(cfg.model.train_ds), model.setup_validation_data(cfg.model.validation_ds) since restored models don't automatically setup their datasets and trainer.

@titu1994 Hello!
When learning from *-lost.ckpt, where does Nemo get the rest of the training parameters from, am I interested in optimal.sched.min_lr?
I trained the model with min_lr = 1.0e-05, then continued training from .ckpt with min_lr = 1.0e-09, but lr did not fall below 1.0e-05.

asr_model = nemo_asr.models.EncDecCTCModelBPE.restore_from(
        cfg.model.init_from_nemo_model.citrinet.path, map_location="cpu"
    ) 
...
...
asr_model.cfg.optim = cfg.model.optim # in config set new value for lr  
trainer = pl.Trainer(**cfg.trainer)  
asr_model.set_trainer(trainer)  
trainer.fit(asr_model)