import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import torchvision
from torchvision import *
from torch.utils.data import Dataset, DataLoader
from nfnets import replace_conv,SGD_AGC

net = models.resnet50(pretrained=True)
replace_conv(net)
net = net.cuda() if device else net
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 13)
net.fc = net.fc.cuda() if use_cuda else net.fc
ct = 0
for name, child in net.named_children():
    ct += 1
    if ct < 8:
        for name2, params in child.named_parameters():
            params.requires_grad = False

criterion = nn.CrossEntropyLoss()
optim_agc = SGD_AGC(net.parameters(), 1e-3,momentum=0.9)

n_epochs = 5
print_every = 10
valid_loss_min = np.Inf
val_loss = []
val_acc = []
train_loss = []
train_acc = []
total_step = len(train_dataloader)

for epoch in range(1, n_epochs+1):
    running_loss = 0.0
    correct = 0
    total=0
    print(f'Epoch {epoch}\n')
    for batch_idx, (data_, target_) in enumerate(train_dataloader):
        data_, target_ = data_.to(device), target_.to(device)
        optim_agc.zero_grad()
        outputs = net(data_)
        loss = criterion(outputs, target_)
        loss.backward()
        optim_agc.step()
        running_loss += loss.item()
        _,pred = torch.max(outputs, dim=1)
        correct += torch.sum(pred==target_).item()
        total += target_.size(0)
        if (batch_idx) % 20 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch, n_epochs, batch_idx, total_step, loss.item()))
    train_acc.append(100 * correct / total)
    train_loss.append(running_loss/total_step)
		

Traceback (most recent call last):

File "image_product_classifier.py", line 74, in optim_agc.step() File "/home/sachin_mohan/venv/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context return func(*args, **kwargs) File "/home/sachin_mohan/venv/lib/python3.6/site-packages/nfnets/sgd_agc.py", line 105, in step grad_norm = unitwise_norm(p.grad) File "/home/sachin_mohan/venv/lib/python3.6/site-packages/nfnets/utils.py", line 27, in unitwise_norm if x.ndim <= 1: AttributeError: 'NoneType' object has no attribute 'ndim'``

I use the code, however, I got the following error.

At first, I got the following;

base.py", line 262, in __init__
    dilation=dilation, groups=groups, bias=bias, padding_mode=padding_mode)
TypeError: __init__() got an unexpected keyword argument 'padding_mode'

When I remove padding_mode=padding_mode, I got another error, as presented below

utils.py", line 23, in replace_conv
    setattr(module, name, torch.nn.Identity())
AttributeError: module 'torch.nn' has no attribute 'Identity'

How to solve this?

Describe the bug Running either replace_conv this code form the readme on the front page: model = vgg16() replace_conv(model, WSConv2d) # This repo's original implementation replace_conv(model, ScaledStdConv2d) # From timm

Results in this error:

  File "/opt/conda/lib/python3.8/site-packages/nfnets/utils.py", line 25, in replace_conv
    replace_conv(mod, conv_class)
  File "/opt/conda/lib/python3.8/site-packages/nfnets/utils.py", line 18, in replace_conv
    setattr(module, name, conv_class(target_mod.in_channels, target_mod.out_channels, target_mod.kernel_size,
  File "/opt/conda/lib/python3.8/site-packages/nfnets/base.py", line 262, in __init__
    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 386, in __init__
    super(Conv2d, self).__init__(
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 108, in __init__
    if bias:
RuntimeError: Boolean value of Tensor with more than one value is ambiguous

Describe the bug Using LambdaLR, it will call len(optimizer.param_groups), but this is not implemented for AGC.

To Reproduce

model = torch.nn.Conv1d(10,20,4)
optimizer = optim.AdamW(model.parameters())
optimizer_agc = AGC(model.parameters(),optimizer)

lambda1 = lambda iteration: 0.05*iteration
scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(optimizer,lr_lambda=lambda1)
scheduler_warmup_agc = torch.optim.lr_scheduler.LambdaLR(optimizer_agc,lr_lambda=lambda1)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-58-f68dd026a6de> in <module>
----> 1 scheduler_warmup2 = torch.optim.lr_scheduler.LambdaLR(optimizer2,lr_lambda=lambda1)

/gpfs/alpine/proj-shared/fus131/conda-envs/torch1.5.0v2/lib/python3.6/site-packages/torch/optim/lr_scheduler.py in __init__(self, optimizer, lr_lambda, last_epoch)
    180
    181         if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple):
--> 182             self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups)
    183         else:
    184             if len(lr_lambda) != len(optimizer.param_groups):

AttributeError: 'AGC' object has no attribute 'param_groups'

Pytorch v1.5.0

Describe the bug A clear and concise description of what the bug is.

To Reproduce Steps to reproduce the behavior:

1. Go to example and instantiate a resnet18
2. Send model to torch.device('cuda)
3. Define a tensor on the gpu
4. Call model.forward()
5. RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/native/cuda/Loops.cuh":94, please report a bug to PyTorch.

Expected behavior Regular output

Screenshots If applicable, add screenshots to help explain your problem.

See here

My Solution:

Its hacky obviously but it works. Simply replace https://github.com/vballoli/nfnets-pytorch/blob/867860eebffcc70fb87a389d770cfd4a73c6b30c/nfnets/base.py#L22 with scale = torch.rsqrt(torch.max(var * fan_in, torch.tensor(eps).to(var.device))) * self.gain.view_as(var).to(var.device)

Hello,

Is there a way to apply AGC externally without modifying the optimizer code?

I am using optimizers from torch_optimizer package and that would be good.

I just did a minor modification on WSConv2d to implement a WS version Conv1d. It works fine with my own model which is a time series model using Transformer + Conv1d, and it removes the overfitting issue in my model. Thanks

Hi @vballoli, It seems to have a bug in the code of AGC:

if model is not None:
    assert ignore_agc not in [None, []], "You must specify ignore_agc for AGC to ignore fc-like(or other) layers"
    names = [name for name, module in model.named_modules()]

    for module_name in ignore_agc:
        if module_name not in names:
            raise ModuleNotFoundError("Module name {} not found in the model".format(module_name))
        params = [{"params": list(module.parameters())} for name,
                          module in model.named_modules() if name not in ignore_agc]
else:
    params = [{"params": list(module.parameters())} for name,
                       module in model.named_modules()]

When model is None then the else part of the code cannot get name and module from model.named_modules().

Thanks

Hello! What's the reason you changed the AGC condition from trigger = grad_norm > max_norm as in original paper and code?

https://github.com/vballoli/nfnets-pytorch/blob/1513cea2c39e09189f4883ad90c2337a8fb9f9ed/nfnets/agc.py#L71

commit: 4b03658f16d5746d45ea27711f5b33551a472e00

resnet18 with replace this layer learn better, but still need two times more epochs for same result than original on CIFAR10

class ScaledStdConv2d(nn.Conv2d):
    """Conv2d layer with Scaled Weight Standardization.
    Paper: `Characterizing signal propagation to close the performance gap in unnormalized ResNets` -
        https://arxiv.org/abs/2101.08692
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=None, dilation=1, groups=1,
                 bias=True, gain=True, gamma=1.0, eps=1e-5, use_layernorm=False):
        if padding is None:
            padding = get_padding(kernel_size, stride, dilation)
        super().__init__(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.gain = nn.Parameter(torch.ones(self.out_channels, 1, 1, 1)) if gain else None
        self.scale = gamma * self.weight[0].numel() ** -0.5  # gamma * 1 / sqrt(fan-in)
        self.eps = eps ** 2 if use_layernorm else eps
        self.use_layernorm = use_layernorm  # experimental, slightly faster/less GPU memory use

    def get_weight(self):
        if self.use_layernorm:
            weight = self.scale * F.layer_norm(self.weight, self.weight.shape[1:], eps=self.eps)
        else:
            std, mean = torch.std_mean(self.weight, dim=[1, 2, 3], keepdim=True, unbiased=False)
            weight = self.scale * (self.weight - mean) / (std + self.eps)
        if self.gain is not None:
            weight = weight * self.gain
        return weight

    def forward(self, x):
        return F.conv2d(x, self.get_weight(), self.bias, self.stride, self.padding, self.dilation, self.groups)

Describe the bug In the paper it is mentioned that AGC is NOT used in the final layer.

Expected behavior It will be great AGC can be disabled for final layer. Adding a check here may be sufficient to get the desired behavior. Please let me know if i am misunderstanding something. Thanks.