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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ema-pytorch
-Version: 0.3.2
+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.2 → 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.2 → ema-pytorch-0.3.3}/ema_pytorch/ema_pytorch.py

@@ -69,6 +69,8 @@ class EMA(Module):
         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
 
         self.include_online_model = include_online_model
@@ -202,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.2 → 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.2
+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.2 → 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.2',
+  version = '0.3.3',
   license='MIT',
   description = 'Easy way to keep track of exponential moving average version of your pytorch module',
   author = 'Phil Wang',