How do we save the finetuned ASR model and use them for inference?

Hello, I was able to fine-tune a model on my dataset as given in the tutorial Pretrained Models and Fine-Tuning. However, I couldn’t find a clear way of saving the fine-tuned model ( something equivalent to torch.save or pickle). I have gone through the checkpoints tutorial and tried to do something as follows (steps after brain.fit() – to save and load the fine-tune model):

from speechbrain.utils.checkpoints import Checkpointer
checkpoint_dir = "./full_example_checkpoints"
checkpointer = Checkpointer(checkpoint_dir)

modules = {"enc": asr_model.mods.encoder.model, 
           "emb": asr_model.hparams.emb,
           "dec": asr_model.hparams.dec,
           "compute_features": asr_model.mods.encoder.compute_features,
           "normalize": asr_model.mods.encoder.normalize,
           "seq_lin": asr_model.hparams.seq_lin, 
}
checkpointer.add_recoverables(modules)
ckpt = checkpointer.save_checkpoint()
ckpt_finder = Checkpointer(checkpoint_dir)
get_ckpt = ckpt_finder.find_checkpoint()

**% load a fresh model and then transfer the parameters of fine-tuned models into that**

new_asr_model = EncoderDecoderASR.from_hparams(source="speechbrain/asr-crdnn-rnnlm-librispeech", savedir="./pretrained_ASR")

**% parameter transfer**
sb.utils.checkpoints.torch_parameter_transfer(new_asr_model.mods.encoder.model, get_ckpt.paramfiles['enc'], device='cpu')
sb.utils.checkpoints.torch_parameter_transfer(new_asr_model.hparams.emb, get_ckpt.paramfiles['emb'], device='cpu')
sb.utils.checkpoints.torch_parameter_transfer(new_asr_model.hparams.dec, get_ckpt.paramfiles['dec'], device='cpu')
sb.utils.checkpoints.torch_parameter_transfer(new_asr_model.mods.encoder.compute_features, best_ckpt.paramfiles['compute_features'], device='cpu')
sb.utils.checkpoints.torch_parameter_transfer(new_asr_model.mods.encoder.normalize, get_ckpt.paramfiles['normalize'], device='cpu')
sb.utils.checkpoints.torch_parameter_transfer(new_asr_model.hparams.seq_lin, get_ckpt.paramfiles['seq_lin'], device='cpu')

I believe there must be a better way to do this. Could you please confirm that above code will work? Also, it would be great if you could provide some hints related to this. Thanks in advance.

[Trikaldarshi]

Hello, Thank you very much for all the help. It is working fine for me now and is very convenient as well. Closing the ticket. Thanks again!.

[Gastron]

Ok so you’ve decided to use the Pretrained interface. Looks ok, the paths provided to your Pretrainer seem a little weird, though I think it would work.

Here’s how I’d change it, though:

1. In your HuggingFace repo, remove CKPT.yaml, brain.ckpt, counter.ckpt, dataloader-TRAIN.ckpt, noam_scheduler.ckpt, and optimizer.ckpt since those are not needed for evaluation.
2. Upload a copy of the tokenizer.ckpt (just in case this changes in the referenced repo).
3. Remove the paths argument (the whole block in your YAML) from Pretrainer - these locations are already the default behaviour, there’s no need to specify them again.
4. Add a README to the HuggingFace repo 😃