save version of package along with model to allow users to downgrade the package if there is an incompatibility

Author: lucidrains

setup.py

@@ -1,5 +1,9 @@
+import sys
 from setuptools import setup, find_packages
 
+sys.path[0:0] = ['stylegan2_pytorch']
+from version import __version__
+
 setup(
   name = 'stylegan2_pytorch',
   packages = find_packages(),
@@ -8,7 +12,7 @@
           'stylegan2_pytorch = stylegan2_pytorch.cli:main',
       ],
   },
-  version = '1.2.4',
+  version = __version__,
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

stylegan2_pytorch/stylegan2_pytorch.py

@@ -27,6 +27,7 @@
 
 import torchvision
 from torchvision import transforms
+from stylegan2_pytorch.version import __version__
 from stylegan2_pytorch.diff_augment import DiffAugment
 
 from pytorch_fid import fid_score
@@ -1164,7 +1165,10 @@ def clear(self):
         self.init_folders()
 
     def save(self, num):
-        save_data = {'GAN': self.GAN.state_dict()}
+        save_data = {
+            'GAN': self.GAN.state_dict(),
+            'version': __version__
+        }
 
         if self.GAN.fp16:
             save_data['amp'] = amp.state_dict()
@@ -1188,11 +1192,14 @@ def load(self, num = -1):
 
         load_data = torch.load(self.model_name(name))
 
-        # make backwards compatible
-        if 'GAN' not in load_data:
-            load_data = {'GAN': load_data}
+        if 'version' in load_data:
+            print(f"loading from version {load_data['version']}")
 
-        self.GAN.load_state_dict(load_data['GAN'])
+        try:
+            self.GAN.load_state_dict(load_data['GAN'])
+        except:
+            print('unable to load save model. please try downgrading the package to the version specified by the saved model')
+            exit()
 
         if self.GAN.fp16 and 'amp' in load_data:
             amp.load_state_dict(load_data['amp'])

stylegan2_pytorch/version.py

@@ -0,0 +1 @@
+__version__ = '1.2.5'