I trained a vanilla vae which I modified from this repository. When I try and use the trained model I am unable to load the weights using load_from_checkpoint. It seems there is a mismatch between my checkpoint object and my lightningModule object.
I have setup an experiment (VAEXperiment) using pytorch-lightning LightningModule. I try to load the weights into the network with:
#building a new model
model = VanillaVAE(**config['model_params'])
model.build_layers()
#loading the weights
experiment = VAEXperiment(model, config['exp_params'])
experiment.load_from_checkpoint(path_to_checkpoint, config['exp_params'])
I also tried:
checkpoint = torch.load(path_to_checkpoint, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
But I get an error
Unexpected key(s) in state_dict: "model.encoder.0.0.weight", "model.encoder.0.0.bias"...
I also followed the issue on https://github.com/PyTorchLightning/pytorch-lightning/issues/924 https://github.com/PyTorchLightning/pytorch-lightning/issues/2798
Why I am I getting this error? Is it because of the encoder and decoder modules in my model? Based on the issues log on git it seems that the error resolved. What am I doing wrong?