How to use pre-trained model's learnable encoder/decoder for Conv-TasNet

[cemox35]

Hello,

I trained the Conv-TasNet using default parameters. I want to use this pre-trained model's encoder and decoder in different Conv-TasNet model(I just change the separation module). Can I do that?

[mpariente]

Yes, do you want them to be fixed or keep learning?

[cemox35]

Thank you for your answer. I want to fix them. How can I do that?

[JorisCos]

I am going to write a small example to show how to proceed, as it is a nice thing to have in your toolbox.

[JorisCos]

if you run this example you can look at the state of the encoder/decoder with model.encoder.state_dict()['filterbank._filters'] and see that nothing changes after training while the weights in model.masker.state_dict() have changed.

[cemox35]

Thank you for your answer. I want to load my model(not from_pretrained(), actually), but I cannot load my auto-checkpoint(in lightning-logs). Could you help that problem? I am very grateful.

[JorisCos]

If you trained with Asteroid you should be able to run ConvTasNet.from_pretrained("best_model.pth")

[JorisCos]

If you just have a .ckpt you could try something like this :

state_dict = torch.load(os.path.join(exp_dir,'your_path_to_your_log.ckpt'))
system.load_state_dict(state_dict=state_dict["state_dict"])
system.cpu()

to_save = system.model.serialize()
torch.save(to_save, os.path.join(exp_dir, "model.pth"))

[cemox35]

Thank you! It works.