openclimatefix / skillful_nowcasting

Implementation of DeepMind's Deep Generative Model of Radar (DGMR) https://arxiv.org/abs/2104.00954
MIT License
211 stars 59 forks source link

Training is fine, but i can not load the trained model #43

Open zhrli opened 1 year ago

zhrli commented 1 year ago

Describe the bug As I run this code, errors reported: import torch from dgmr import DGMR, Sampler, Generator, Discriminator, LatentConditioningStack, ContextConditioningStack

model = DGMR( forecast_steps=20, #20 input_channels=1, output_shape=256, latent_channels=384, context_channels=192, num_samples=3, )

model.sampler = Sampler.from_pretrained("openclimatefix/dgmr-sampler") model.sampler.forecast_steps = 20 model.sampler.output_shape = 256 model.discriminator = Discriminator.from_pretrained("openclimatefix/dgmr-discriminator") model.latent_stack = LatentConditioningStack.from_pretrained("openclimatefix/dgmr-latent-conditioning-stack") model.context_stack = ContextConditioningStack.from_pretrained("openclimatefix/dgmr-context-conditioning-stack") model.generator = Generator(conditioning_stack=model.context_stack, latent_stack=model.latent_stack, sampler=model.sampler) model=DGMR.load_from_checkpoint(checkpoint_path='best-v3.ckpt', strict=False)
print(model.config)

To Reproduce Steps to reproduce the behavior: RuntimeError: Error(s) in loading state_dict for DGMR: size mismatch for latent_stack.l_block1.conv_1x1.weight: copying a param with shape torch.Size([16, 8, 1, 1]) from checkpoint, the shape in current model is torch.Size([4, 8, 1, 1]). ... ... ... Process finished with exit code 1

Expected behavior Training is fine . but I don't know why i can not load the trained model.

Additional context Add any other context about the problem here.