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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ema-pytorch
-Version: 0.3.1
+Version: 0.3.2
 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.2}/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,8 @@ 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
 
         # whether to include the online model within the module tree, so that state_dict also saves it
 
@@ -127,6 +129,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()

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ema-pytorch
-Version: 0.3.1
+Version: 0.3.2
 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.2}/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.2',
   license='MIT',
   description = 'Easy way to keep track of exponential moving average version of your pytorch module',
   author = 'Phil Wang',