sophia v2 and import error

Sophia/Sophiav2.py

@@ -0,0 +1,67 @@
+import torch
+
+class Sophia2(torch.optim.Optimizer):
+    def __init__(self, model, input_data, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, k=10, estimator="Hutchinson", rho=1):
+        self.model = model
+        self.input_data = input_data
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, k=k, estimator=estimator, rho=rho)
+        super(Sophia2, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            eps = group['eps']
+            rho = group['rho']
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Sophia does not support sparse gradients")
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['m'] = torch.zeros_like(p.data)
+                    state['h'] = torch.zeros_like(p.data)
+
+                m, h = state['m'], state['h']
+                beta1, beta2 = group['betas']
+                state['step'] += 1
+
+                if group['weight_decay'] != 0:
+                    grad = grad.add(group["weight_decay"], p.data)
+
+                m.mul_(beta1).add_(1 - beta1, grad)
+
+                if state['step'] % group['k'] == 1:
+                    if group['estimator'] == "Hutchinson":
+                        hessian_estimate = self.hutchinson(p, grad)
+                    elif group['estimator'] == "Gauss-Newton-Bartlett":
+                        hessian_estimate = self.gauss_newton_bartlett(p, grad)
+                    else:
+                        raise ValueError("Invalid estimator choice")
+                    h.mul_(beta2).add_(1 - beta2, hessian_estimate)
+
+                with torch.no_grad():
+                    p.data.add_(-group['lr'] * group['weight_decay'], p.data)
+                    p.data.addcdiv_(-group['lr'], m, h.add(eps).clamp(max=rho))
+
+        return loss
+
+    def hutchinson(self, p, grad):
+        u = torch.randn_like(grad)
+        grad_dot_u = torch.einsum("...,...->", grad, u)
+        hessian_vector_product = torch.autograd.grad(grad_dot_u, p, retain_graph=True)[0]
+        return u * hessian_vector_product
+
+    def gauss_newton_bartlett(self, p, grad):
+        B = len(self.input_data)
+        logits = [self.model(xb) for xb in self.input_data]
+        y_hats = [torch.softmax(logit, dim=0) for logit in logits]
+        g_hat = torch.autograd.grad(sum([self.loss_function(logit, y_hat) for logit, y_hat in zip(logits, y_hats)]) / B, p, retain_graph=True)[0]
+        return B * g_hat * g_hat

Sophia/__init__.py

@@ -1,3 +1,6 @@
 from Sophia.Sophia import Sophia
-from decoupled_sophia.decoupled_sophia.decoupled_sophia import DecoupledSophia
-from experiments.training import trainer
+from Sophia.Sophiav2 import Sophia2 
+# from decoupled_sophia.decoupled_sophia.decoupled_sophia import DecoupledSophia
+from experiments.training import trainer
+
+from Sophia.decoupled_sophia.decoupled_sophia import DecoupledSophia

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'Sophia-Optimizer',
   packages = find_packages(exclude=[]),
-  version = '0.1.6',
+  version = '0.1.8',
   license='APACHE',
   description = 'Sophia Optimizer ULTRA FAST',
   author = 'Kye Gomez',