ema-pytorch 0.3.0__py3-none-any.whl → 0.3.2__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,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
 
@@ -83,7 +85,8 @@ class EMA(Module):
         if not exists(self.ema_model):
             try:
                 self.ema_model = deepcopy(model)
-            except:
+            except Exception as e:
+                print(f'Error: While trying to deepcopy model: {e}')
                 print('Your model was not copyable. Please make sure you are not using any LazyLinear')
                 exit()
 
@@ -126,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.0.dist-info → ema_pytorch-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ema-pytorch
-Version: 0.3.0
+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.2.dist-info/RECORD

@@ -0,0 +1,7 @@
+ema_pytorch/__init__.py,sha256=jCgUV6FfA65ct6lMEClKKwl-gjZMDv7aOeaN3RshE0U,40
+ema_pytorch/ema_pytorch.py,sha256=WBOhcK_7WWBy48tYcsIJceaZvF4oz8qjOueQX1LoLqI,8489
+ema_pytorch-0.3.2.dist-info/LICENSE,sha256=xZDkKtpHE2TPCAeqKe1fjdpKernl1YW-d01j_1ltkAU,1066
+ema_pytorch-0.3.2.dist-info/METADATA,sha256=uVi-SjuBOrGBWKLct4GHX_0x16eeKVIp7795jRbaDp4,715
+ema_pytorch-0.3.2.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+ema_pytorch-0.3.2.dist-info/top_level.txt,sha256=XXFJmHviark_32Hfm5X9niezVmnRTUIhfdifCrJgXmE,12
+ema_pytorch-0.3.2.dist-info/RECORD,,

{ema_pytorch-0.3.0.dist-info → ema_pytorch-0.3.2.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.0.dist-info/RECORD

@@ -1,7 +0,0 @@
-ema_pytorch/__init__.py,sha256=jCgUV6FfA65ct6lMEClKKwl-gjZMDv7aOeaN3RshE0U,40
-ema_pytorch/ema_pytorch.py,sha256=4OPFewDiKVezIvEdvXvn4d4PGrXafhlH39qg_5lrMGk,8102
-ema_pytorch-0.3.0.dist-info/LICENSE,sha256=xZDkKtpHE2TPCAeqKe1fjdpKernl1YW-d01j_1ltkAU,1066
-ema_pytorch-0.3.0.dist-info/METADATA,sha256=uzoXWFnqKpq9tFZGU1oK1ZKJSEqZ_CFiMmAy2LLYt4A,715
-ema_pytorch-0.3.0.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-ema_pytorch-0.3.0.dist-info/top_level.txt,sha256=XXFJmHviark_32Hfm5X9niezVmnRTUIhfdifCrJgXmE,12
-ema_pytorch-0.3.0.dist-info/RECORD,,