search

nickyisadog / latent-diffusion-inpainting

80 stars 14 forks source link

[Fix] Modify Weight Storage to Include Additional Information #21

Closed gotjd709 closed 3 months ago

gotjd709 commented 3 months ago

The method you suggested is excellent and has greatly assisted my work. However, while proceeding with step 2 of the Usage, an error occurred during model training because only the state_dict was saved when storing the model weights. I understand that information such as 'epoch' and 'global_step' also needs to be saved in the weights for training with PyTorch Lightning. Therefore, I have modified the code, and this resolved the error that was not resolved in my case. I am submitting this pull request to inform others about this issue and solution. Thank you.

ldm_model = torch.load(ldm_model_path)
torch.save({
            'epoch': ldm_model['epoch'],
            'global_step': ldm_model['global_step'],
            'pytorch-lightning_version': ldm_model['pytorch-lightning_version'],
            'state_dict': model.state_dict(),
            'callbacks': ldm_model['callbacks'],
            'optimizer_states': ldm_model['optimizer_states'],
            'lr_schedulers': ldm_model['lr_schedulers'],
            }, 'updated_ldm.ckpt')