Doubt about resuming training

[xv5kVu4FN]

Hi! First of all, thanks for sharing this awesome repository.
I have a doubt regarding how to resume training.

In the code, the diffusion_process_trainer calls the load method, where the ckpt/pretrained.pt is loaded and we do: self.step = data['step']
Now, I was peeking on the values that pretrained.pt has and I see that step is set to zero.

If I now run train.py, will it use pretrained.pt and resume the the training or will it start from scratch?

Thanks!

[Baixuzx7]

Hi! Thank you for your kind words. We set the step value to 0 in the pretrained checkpoint to prevent any potential bugs and to make it easier for fine-tuning. Therefore, it's normal for the step value to be 0 when loading the checkpoint.

When you run train.py after loading pretrained.pt, it will correctly resume the training from the point where the checkpoint was saved. It won't start from scratch.