ema-pytorch 0.3.1__tar.gz → 0.3.3__tar.gz

{ema-pytorch-0.3.1 → ema-pytorch-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ema-pytorch
-Version: 0.3.1
+Version: 0.3.3
 Summary: Easy way to keep track of exponential moving average version of your pytorch module
 Home-page: https://github.com/lucidrains/ema-pytorch
 Author: Phil Wang

{ema-pytorch-0.3.1 → ema-pytorch-0.3.3}/README.md

@@ -52,3 +52,16 @@ ema_output = ema(data)
 ## Todo
 
 - [ ] address the issue of annealing EMA to 1 near the end of training for BYOL https://github.com/lucidrains/byol-pytorch/issues/82
+
+## Citations
+
+```bibtex
+@article{Karras2023AnalyzingAI,
+    title   = {Analyzing and Improving the Training Dynamics of Diffusion Models},
+    author  = {Tero Karras and Miika Aittala and Jaakko Lehtinen and Janne Hellsten and Timo Aila and Samuli Laine},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2312.02696},
+    url     = {https://api.semanticscholar.org/CorpusID:265659032}
+}
+```

{ema-pytorch-0.3.1 → ema-pytorch-0.3.3}/ema_pytorch/ema_pytorch.py

@@ -43,7 +43,7 @@ class EMA(Module):
 
     Args:
         inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
-        power (float): Exponential factor of EMA warmup. Default: 1.
+        power (float): Exponential factor of EMA warmup. Default: 2/3.
         min_value (float): The minimum EMA decay rate. Default: 0.
     """
 
@@ -53,6 +53,7 @@ class EMA(Module):
         model: Module,
         ema_model: Optional[Module] = None,           # if your model has lazylinears or other types of non-deepcopyable modules, you can pass in your own ema model
         beta = 0.9999,
+        karras_beta = False,                          # if True, uses the karras time dependent beta
         update_after_step = 100,
         update_every = 10,
         inv_gamma = 1.0,
@@ -65,7 +66,10 @@ class EMA(Module):
         allow_different_devices = False               # if the EMA model is on a different device (say CPU), automatically move the tensor
     ):
         super().__init__()
-        self.beta = beta
+        self._beta = beta
+        self.karras_beta = karras_beta
+
+        self.is_frozen = beta == 1.
 
         # whether to include the online model within the module tree, so that state_dict also saves it
 
@@ -127,6 +131,13 @@ class EMA(Module):
     @property
     def model(self):
         return self.online_model if self.include_online_model else self.online_model[0]
+    
+    @property
+    def beta(self):
+        if self.karras_beta:
+            return (1 - 1 / (self.step + 1)) ** (1 + self.power)
+
+        return self._beta
 
     def eval(self):
         return self.ema_model.eval()
@@ -193,6 +204,9 @@ class EMA(Module):
 
     @torch.no_grad()
     def update_moving_average(self, ma_model, current_model):
+        if self.is_frozen:
+            return
+
         copy, lerp = self.inplace_copy, self.inplace_lerp
         current_decay = self.get_current_decay()
 

{ema-pytorch-0.3.1 → ema-pytorch-0.3.3}/ema_pytorch.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ema-pytorch
-Version: 0.3.1
+Version: 0.3.3
 Summary: Easy way to keep track of exponential moving average version of your pytorch module
 Home-page: https://github.com/lucidrains/ema-pytorch
 Author: Phil Wang

{ema-pytorch-0.3.1 → ema-pytorch-0.3.3}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'ema-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.3.1',
+  version = '0.3.3',
   license='MIT',
   description = 'Easy way to keep track of exponential moving average version of your pytorch module',
   author = 'Phil Wang',