Make AdaHessian & SophiaH functional

Author: ferris

pytorch_optimizer/__init__.py

@@ -83,6 +83,8 @@
 from pytorch_optimizer.optimizer.sm3 import SM3
 from pytorch_optimizer.optimizer.srmm import SRMM
 from pytorch_optimizer.optimizer.swats import SWATS
+from pytorch_optimizer.optimizer.adahessian import AdaHessian
+from pytorch_optimizer.optimizer.sophiah import SophiaH
 from pytorch_optimizer.optimizer.utils import (
     clip_grad_norm,
     disable_running_stats,
@@ -147,6 +149,8 @@
     AdaShift,
     AdaDelta,
     Amos,
+    AdaHessian,
+    SophiaH
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/base/optimizer.py

@@ -34,9 +34,11 @@ def compute_hutchinson_hessian(self, nsamples: int = 1, pre_zero=True, alpha=1.0
         grads = [p.grad for p in params]
 
         for i in range(nsamples):
+            # Gaussian N(0,Id)
+            zs = [torch.randn(p.size(), device=p.device) for p in params]
             # Rademacher distribution {-1.0, 1.0}
-            zs = [torch.randint(0, 2, p.size(), device=p.device) * 2.0 - 1.0 for p in params]
-            h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, only_inputs=True, retain_graph=i < nsamples - 1)
+            # zs = [torch.randint(0, 2, p.size(), device=p.device) * 2.0 - 1.0 for p in params]
+            h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, retain_graph=i < nsamples - 1)
             for h_z, z, p in zip(h_zs, zs, params):
                 # approximate the expected values of z*(H@z)
                 self.state[p]['hessian'].add_(h_z * z, alpha=1/nsamples * alpha)

pytorch_optimizer/optimizer/adahessian.py

@@ -43,20 +43,29 @@ def __init__(self,
         self.validate_non_negative(eps, 'eps')
         self.validate_range(hessian_power, "Hessian Power", 0, 1, range_type='(]')
 
+        self.update_period = update_period
+        self.n_samples = n_samples
         defaults: DEFAULTS = {
             'lr': lr,
             'betas': betas,
             'weight_decay': weight_decay,
             'weight_decouple': weight_decouple,
             'fixed_decay': fixed_decay,
             'hessian_power': hessian_power,
-            'update_period': update_period,
-            'n_samples': n_samples,
             'eps': eps,
         }
         self._step = 0
         super().__init__(params, defaults)
 
+    @torch.no_grad()
+    def reset(self):
+        self._step = 0
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_hessian_diag_sq'] = torch.zeros_like(p)
+
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
         loss: LOSS = None
@@ -72,9 +81,6 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if p.grad is None:
                     continue
 
-                if self.average_conv_kernel and p.dim() == 4:
-                    p.hess = torch.abs(p.hess).mean(dim=[2, 3], keepdim=True).expand_as(p.hess).clone()
-
                 grad = p.grad
                 if grad.is_sparse:
                     raise NoSparseGradientError(str(self))
@@ -100,10 +106,12 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 # Decay the first and second moment running average coefficient
                 exp_avg.mul_(beta1).add_(p.grad, alpha=1 - beta1)
                 if self._step % self.update_period == 0:
+                    # if self.average_conv_kernel and p.dim() == 4:
+                    #     state['hessian'] = torch.abs(state['hessian']).mean(dim=[2, 3], keepdim=True).expand_as(state['hessian']).clone()
                     exp_hessian_diag_sq.mul_(beta2).addcmul_(state['hessian'], state['hessian'], value=1 - beta2)
 
-                bias_correction1 = 1 - beta1 ** self._step
-                bias_correction2 = 1 - beta2 ** self._step
+                bias_correction1 = 1 - beta1 ** (self._step+1)
+                bias_correction2 = 1 - beta2 ** (self._step+1)
 
                 k = group['hessian_power']
                 denom = (exp_hessian_diag_sq / bias_correction2).pow_(k / 2).add_(group['eps'])

pytorch_optimizer/optimizer/sophiah.py

@@ -48,13 +48,22 @@ def __init__(self,
             'weight_decouple': weight_decouple,
             'fixed_decay': fixed_decay,
             'p': p,
-            'update_period': update_period,
-            'n_samples': n_samples,
             'eps': eps,
         }
+        self.n_samples = n_samples
+        self.update_period = update_period
         self._step = 0
         super().__init__(params, defaults)
 
+    @torch.no_grad()
+    def reset(self):
+        self._step = 0
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                state['momentum'] = torch.zeros_like(p)
+                state['hessian_moment'] = torch.zeros_like(p)
+
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
         loss: LOSS = None
@@ -63,7 +72,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 loss = closure()
 
         if self._step % self.update_period == 0:
-            self.compute_hutchinson_hessian(self.n_smaples)
+            self.compute_hutchinson_hessian(self.n_samples)
 
         for group in self.param_groups:
             for p in group['params']:
@@ -77,8 +86,8 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 # State initialization
                 state = self.state[p]
                 if 'momentum' not in state:
-                    state['momentum'] = torch.zeros_like(p.data)
-                    state['hessian_moment'] = torch.zeros_like(p.data)
+                    state['momentum'] = torch.zeros_like(p)
+                    state['hessian_moment'] = torch.zeros_like(p)
 
                 self.apply_weight_decay(
                     p=p,
@@ -100,7 +109,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 # The official implementation uses a different method to achieve the same thing (might be faster?):
                 # https://github.com/Liuhong99/Sophia/blob/bff9df9b584e2084fe037af1ab38f4db31f0acca/sophia.py#L201
                 update = torch.clip(momentum/torch.clip(hessian_moment, group['eps']), -group['p'], group['p'])
-                p.add_(update, value=-group['lr'])
+                p.add_(update, alpha=-group['lr'])
 
         self._step += 1
         return loss