start saving version of library in the checkpoints, so in case of breaking changes, researchers can return to the right working version

Author: lucidrains

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -23,6 +23,8 @@
 
 from accelerate import Accelerator
 
+from denoising_diffusion_pytorch.version import __version__
+
 # constants
 
 ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
@@ -828,7 +830,8 @@ def save(self, milestone):
             'model': self.accelerator.get_state_dict(self.model),
             'opt': self.opt.state_dict(),
             'ema': self.ema.state_dict(),
-            'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None
+            'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None,
+            'version': __version__
         }
 
         torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
@@ -846,6 +849,9 @@ def load(self, milestone):
         self.opt.load_state_dict(data['opt'])
         self.ema.load_state_dict(data['ema'])
 
+        if 'version' in data:
+            print(f"loading from version {data['version']}")
+
         if exists(self.accelerator.scaler) and exists(data['scaler']):
             self.accelerator.scaler.load_state_dict(data['scaler'])
 

denoising_diffusion_pytorch/version.py

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

setup.py

@@ -1,9 +1,11 @@
 from setuptools import setup, find_packages
 
+exec(open('denoising_diffusion_pytorch/version.py').read())
+
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.32.0',
+  version = __version__,
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',