How to save and load LightningModule whose input containing the pretrained moduel?
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?
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)