safe division

Author: lucidrains

hamburger_pytorch/hamburger_pytorch.py

@@ -24,7 +24,8 @@ def __init__(
         dim,
         n,
         ratio = 8,
-        K = 6
+        K = 6,
+        eps = 2e-8
     ):
         super().__init__()
         r = dim // ratio
@@ -36,8 +37,10 @@ def __init__(
         self.D = nn.Parameter(D)
         self.C = nn.Parameter(C)
 
+        self.eps = eps
+
     def forward(self, x):
-        b, D, C = x.shape[0], self.D, self.C
+        b, D, C, eps = x.shape[0], self.D, self.C, self.eps
 
         # x is made non-negative with relu as proposed in paper
         x = F.relu(x)
@@ -52,8 +55,8 @@ def forward(self, x):
             # only calculate gradients on the last step, per propose 'One-step Gradient'
             context = null_context if k == 0 else torch.no_grad
             with context():
-                C_new = C * ((t(D) @ x) / (t(D) @ D @ C))
-                D_new = D * ((x @ t(C)) / (D @ C @ t(C)))
+                C_new = C * ((t(D) @ x) / ((t(D) @ D @ C) + eps))
+                D_new = D * ((x @ t(C)) / ((D @ C @ t(C)) + eps))
                 C, D = C_new, D_new
 
         return D @ C

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'hamburger-pytorch',
   packages = find_packages(),
-  version = '0.0.2',
+  version = '0.0.3',
   license='MIT',
   description = 'Hamburger - Pytorch',
   author = 'Phil Wang',