ema-pytorch 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

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.dist-info → ema_pytorch-0.3.3.dist-info}/METADATA

@@ -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.3.dist-info/RECORD

@@ -0,0 +1,7 @@
+ema_pytorch/__init__.py,sha256=jCgUV6FfA65ct6lMEClKKwl-gjZMDv7aOeaN3RshE0U,40
+ema_pytorch/ema_pytorch.py,sha256=eOycxsyO3aCBzToiBGEqOJIGxklX3N3Yitq3wldwHkg,8573
+ema_pytorch-0.3.3.dist-info/LICENSE,sha256=xZDkKtpHE2TPCAeqKe1fjdpKernl1YW-d01j_1ltkAU,1066
+ema_pytorch-0.3.3.dist-info/METADATA,sha256=i2-cAyNXNi4l-xUVsrvFP5wpJHvzyglB1BBBFhobwqg,715
+ema_pytorch-0.3.3.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+ema_pytorch-0.3.3.dist-info/top_level.txt,sha256=XXFJmHviark_32Hfm5X9niezVmnRTUIhfdifCrJgXmE,12
+ema_pytorch-0.3.3.dist-info/RECORD,,

{ema_pytorch-0.3.1.dist-info → ema_pytorch-0.3.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.3)
+Generator: bdist_wheel (0.42.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ema_pytorch-0.3.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-ema_pytorch/__init__.py,sha256=jCgUV6FfA65ct6lMEClKKwl-gjZMDv7aOeaN3RshE0U,40
-ema_pytorch/ema_pytorch.py,sha256=ALWUc6STWA7ZqACn4OuXAiWqb57tk3Vj5B-RowHKNYU,8186
-ema_pytorch-0.3.1.dist-info/LICENSE,sha256=xZDkKtpHE2TPCAeqKe1fjdpKernl1YW-d01j_1ltkAU,1066
-ema_pytorch-0.3.1.dist-info/METADATA,sha256=dqsiNPoAH7aG_Na--NPg4P_DsgbpJjcc3KsA1lu6e-I,715
-ema_pytorch-0.3.1.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-ema_pytorch-0.3.1.dist-info/top_level.txt,sha256=XXFJmHviark_32Hfm5X9niezVmnRTUIhfdifCrJgXmE,12
-ema_pytorch-0.3.1.dist-info/RECORD,,