Publishing weights for fine tuning

[xavriley]

Hi - firstly, thank you for this wonderful work. The results I've had so far are fantastic.

I'm interested in fine tuning some of the published checkpoints with some data augmentation strategies to make them a little bit more robust to different recording conditions. Unfortunately the published weights are "weights only" at the moment. For example, loading the published checkpoint I get:

>>> weights_only.keys()
dict_keys(['state_dict'])

whereas training a model from scratch gives me

>>> full_model.keys()
dict_keys(['state_dict', 'best_state_dict', 'epoch', 'nIter', 'loss_tracker', 'optimizer_state_dict', 'lr_scheduler_state_dict'])

Would it be possible to share the full model from the paper somewhere?

[Yujia-Yan]

Other parts are useless for finetuning as you definitely need to reinitialize all of them. A simple solution would be to load the weights after initializing the checkpoint.

Transkun/transkun/train.py

Line 65 in c5cb837