File size: 3,365 Bytes
998bb30 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
# Code is copied from
# github.com/davda54/sam/sam.py@3c3afdb
import torch
from torch.nn.modules.batchnorm import _BatchNorm
class SAM(torch.optim.Optimizer):
def __init__(self,
params,
base_optimizer,
rho=0.05,
adaptive=False,
**kwargs):
assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
super(SAM, self).__init__(params, defaults)
self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
self.param_groups = self.base_optimizer.param_groups
self.defaults.update(self.base_optimizer.defaults)
@torch.no_grad()
def first_step(self, zero_grad=False):
grad_norm = self._grad_norm()
for group in self.param_groups:
scale = group['rho'] / (grad_norm + 1e-12)
for p in group['params']:
if p.grad is None:
continue
self.state[p]['old_p'] = p.data.clone()
e_w = (torch.pow(p, 2)
if group['adaptive'] else 1.0) * p.grad * scale.to(p)
p.add_(e_w) # climb to the local maximum "w + e(w)"
if zero_grad:
self.zero_grad()
@torch.no_grad()
def second_step(self, zero_grad=False):
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
p.data = self.state[p][
'old_p'] # get back to "w" from "w + e(w)"
self.base_optimizer.step() # do the actual "sharpness-aware" update
if zero_grad:
self.zero_grad()
@torch.no_grad()
def step(self, closure=None):
assert closure is not None, \
'Sharpness Aware Minimization requires closure, ' \
'but it was not provided'
closure = torch.enable_grad()(
closure) # the closure should do a full forward-backward pass
self.first_step(zero_grad=True)
closure()
self.second_step()
def _grad_norm(self):
# put everything on the same device, in case of model parallelism
shared_device = self.param_groups[0]['params'][0].device
norm = torch.norm(torch.stack([
((torch.abs(p) if group['adaptive'] else 1.0) *
p.grad).norm(p=2).to(shared_device) for group in self.param_groups
for p in group['params'] if p.grad is not None
]),
p=2)
return norm
def load_state_dict(self, state_dict):
super().load_state_dict(state_dict)
self.base_optimizer.param_groups = self.param_groups
def disable_running_stats(model):
def _disable(module):
if isinstance(module, _BatchNorm):
module.backup_momentum = module.momentum
module.momentum = 0
model.apply(_disable)
def enable_running_stats(model):
def _enable(module):
if isinstance(module, _BatchNorm) and hasattr(module,
'backup_momentum'):
module.momentum = module.backup_momentum
model.apply(_enable)
|