How to save and load LightningModule whose input containing the pretrained moduel?

[vrvrv]

Hi,

I'm applying Pytorch Lightning module to VAE and our model We first train VAE and give the best checkpoint of pretrained VAE as the initial weight of our model.

# STEP 1. Train VAE
vae = VAE(...)

trainer = Trainer(...)
trainer.fit(vae)

# STEP 2.
vae = VAE.load_from_checkpoint(...)

class Model(LightningModule):
    def __init__(self, encoder, decoder, learning_rate):
       super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        
        self.save_hyperparameters("learning_rate")
        ...

encoder = copy.deepcopy(vae.encoder)
decoder = copy.deepcopy(vae.decoder)

model = Model(
    encoder=encoder,
    decoder=decoder,
    ...
)
trainer.fit(model)

The problem is when I load the model after train ends. Since the torch modules are contained in input arguments of Model, the common approach

model = Model.load_from_checkpoint(...) 

yields following error messages. TypeError: __init__() missing 2 required positional arguments: 'encoder' and 'decoder'
So, what is the best practice for saving and loading the model which uses the pre-trained model?

[rohitgr7]

since part of your model is inside arguments, you can randomly initialize your VAE and let the new checkpoint configure it's weights

vae = VAE()
encoder = copy.deepcopy(vae.encoder)
decoder = copy.deepcopy(vae.decoder)

model = Model.load_from_checkpoint(..., encoder=encoder, decoder=decoder) 

[vrvrv]

Thanks @rohitgr7 ! But, the problem is that during training, I update the parameters of vae.decoder in model while vae.encoder being fixed. Would it be okay even if I use the above code?

[rohitgr7]

yes, in that case while training Model if you have created checkpoints, then Model.load_from_checkpoint(ckpt_path, ...), ckpt_path should contain the freezed pre-trained weights of vae.encoder.

[vrvrv]

Thanks a lot! It works :)