samsunglabs/asam

Implementation of ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks, ICML 2021.

asam's Issues

Multi-GPU training

Hi, thanks for your open-source implementation of ASAM. I am wondering whether if ASAM support multi-GPU training, and how to calculate grad norm when using multiple GPUs? Thanks again and looking forward to your reply.

Code for training imagenet

PyramidNet-272 result reproduce

Hi
Your published code does not have PyramidNet, isn't it? It seems timm has no PyramidNet too. I need to replicate your result with that model since it shows the best result.

Thanks

How to visualize ResNet loss landscape

How to visualize ResNet loss landscape

T_w is applied twice?

I was wondering why is the gradient multiplied by $T_w$ twice in ASAM, p.grad.mul_(t_w)?

ASAM/asam.py

Lines 25 to 33 in f156a68

	p.grad.mul_(t_w)
	wgrads.append(torch.norm(p.grad, p=2))
	wgrad_norm = torch.norm(torch.stack(wgrads), p=2) + 1.e-16
	for n, p in self.model.named_parameters():
	if p.grad is None:
	continue
	t_w = self.state[p].get("eps")
	if 'weight' in n:
	p.grad.mul_(t_w)