Save and overwrite best checkpoint every epoch

[cvillela]

Hello,

Congratulations on the model! I have been experimenting with model training and finetuning ever since the code got released. One feature that could be immensely useful is the writing/overwriting of the best state of the current training at every validation epoch by a given metric (cross-entropy as default). From what I understand, the best state is recognized and saved only at the end of the training, which can be tricky for models that eventually overfit or need to have the training abruptly stopped.

If there is a config parameter for this already, could anyone please point it to me?

Best, Caio

[adefossez]

This is already happening. All the evals happen with the best checkpoint. By default Cross entropy on the valid set is used to select the best checkpoint: https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/solvers/musicgen.py#L46

[adefossez]

The best state is stored in the latest checkpoint of rank 0 (the one without any integer suffix after the .th), under the key best_state or fsdp_best_state (if trained with fsdp). This is the model that is exported when following those isntructions: https://github.com/facebookresearch/audiocraft/blob/main/docs/MUSICGEN.md#importing--exporting-models

[cvillela]

@adefossez Thank you for the response. I was under the impression that the checkpoint of rank 0 was always saving the last state, and was not aware of the best_state key.

I am marking this issue as closed!