Штраф за градиент активации

Я думаю, что ваш код в конечном итоге вычисляет некоторые градиенты дважды на каждом шаге. Я также подозреваю, что он на самом деле никогда не обнуляет градиенты активации, поэтому они накапливаются по шагам.

В целом:

• x.backward() вычисляет градиент x относительно. листья графа вычислений (например, тензоры весов и другие переменные), а также относительно. узлы, явно отмеченные retain_grad(). Он накапливает вычисленный градиент в атрибутах .grad тензоров.

• autograd.grad(x, [y, z]) возвращает градиент x относительно. y и z независимо от того, сохранят ли они градус или нет. По умолчанию он также накапливает градиент во всех атрибутах .grad листьев. Вы можете предотвратить это, передав only_inputs=True.

Я предпочитаю использовать backward() только для шага оптимизации и autograd.grad() всякий раз, когда моей целью является получение «овеществленных» градиентов в качестве промежуточных значений для другого вычисления. Таким образом, я могу быть уверен, что никакие нежелательные градиенты не останутся в атрибутах .grad тензоров после того, как я закончу с ними.

import torch
from torch import nn
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=5)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5)
        self.pool = nn.MaxPool2d(2, 2)
        self.relu = nn.ReLU()
        self.linear = nn.Linear(64 * 5 * 5, 10)

    def forward(self, input):
        conv1 = self.conv1(input)
        pool1 = self.pool(conv1)
        self.relu1 = self.relu(pool1)
        conv2 = self.conv2(self.relu1)
        pool2 = self.pool(conv2)
        self.relu2 = self.relu(pool2)
        relu2 = self.relu2.view(self.relu2.size(0), -1)
        return self.linear(relu2)


model = Net()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
grad_penalty_weight = 10.

for i in range(1000000):
    # Random input and labels; we're not really learning anything
    input = torch.rand(1, 3, 32, 32)
    label = torch.randint(0, 10, (1,))

    output = model(input)
    loss = nn.CrossEntropyLoss()(output, label)

    # This is where the activation gradients are computed
    # only_inputs is optional here, since we're going to call optimizer.zero_grad() later
    # But it makes clear that we're *only* interested in the activation gradients at this point
    grads = torch.autograd.grad(loss, [model.relu1, model.relu2], create_graph=True, only_inputs=True)

    grad_norm = 0
    for grad in grads:
        grad_norm += grad.pow(2).sum()

    optimizer.zero_grad()
    loss = loss + grad_norm * grad_penalty_weight
    loss.backward()
    optimizer.step()

Этот код похоже работает, поскольку градиенты активации становятся меньше. Я не могу комментировать жизнеспособность этого метода как метода регуляризации.