``torch.autograd.grad" 的内部变量是否可微分?

Is ``torch.autograd.grad" differentiable for its internal variables?

我有一个包含可学习梯度下降步骤的网络,实现如下:


实施#1:

import torch
from torch import nn

def cal_grad(x, loss):
    v = torch.autograd.Variable(x, requires_grad=True)
    loss_v = loss(v)
    return torch.autograd.grad(loss_v, v, torch.ones_like(loss_v))[0]

class Net_autograd(nn.Module):
    def __init__(self):
        super(Net_autograd, self).__init__()
        self.A = nn.Parameter(torch.rand(100, 50))
        self.rho = nn.Parameter(torch.rand(1,))

    def forward(self, x_init, y):
        grad = cal_grad(x_init, lambda z: 0.5 * (z.mm(self.A) - y).pow(2).sum(dim=[1,]))  # calculate gradient implicitly (without giving its formula)
        x_hat = x_init - self.rho * grad
        return x_hat

net = Net_autograd()
x_init = torch.rand(10, 100, requires_grad=True)
y = torch.rand(10, 50, requires_grad=False)
x_gt = torch.rand(10, 100, requires_grad=False)

x_hat = net(x_init, y)

loss = (x_hat - x_gt).pow(2).mean()  # L2 loss
loss.backward()

print('A', net.A.grad)
print('rho', net.rho.grad)
print('x_init', x_init.grad)

在它的转发过程中,我用一个可学习的矩阵self.A来计算梯度并得到x_hat,这并不常见。我很困惑为什么在 loss.backward() 之后,参数 self.A 的梯度是 None?

并且我通过显式给出解析梯度计算公式重新实现了原始网络,参数self.A可以成功学习,因为net.A.grad不会是None:


实施#2:

import torch
from torch import nn

class Net_explicit(nn.Module):
    def __init__(self):
        super(Net_explicit, self).__init__()
        self.A = nn.Parameter(torch.rand(100, 50))
        self.rho = nn.Parameter(torch.rand(1,))

    def forward(self, x_init, y):
        grad = (x_init.mm(self.A) - y).mm(self.A.t())  # calculate gradient explicitly
        x_hat = x_init - self.rho * grad
        return x_hat

net = Net_explicit()
x_init = torch.rand(10, 100, requires_grad=True)
y = torch.rand(10, 50, requires_grad=False)
x_gt = torch.rand(10, 100, requires_grad=False)

x_hat = net(x_init, y)

loss = (x_hat - x_gt).pow(2).mean()  # L2 loss
loss.backward()

print('A', net.A.grad)
print('rho', net.rho.grad)
print('x_init', x_init.grad)

我想问的是:

  1. 第二个实现如我所料运行良好,但为什么第一个不训练 self.A(具有 None 梯度)?我猜 torch.autograd.grad 本身可能无法为其内部变量区分,因为外部可学习梯度下降步长 self.rho 和输入 x_init 获得了它们自己的 grads.

  2. 要实现更复杂的管道,其中self.A甚至可能是一个大网络,我该怎么做才能实现#1#2 一样有效(即,使 self.A 可以使用 autograd 机制进行训练)?

  3. 顺便说一句,我不确定x_initgrad是否在#1和[的转发中计算出来=59=]#2相同或完全相等。你能帮我检查一下吗?

注:x_inityx_gt的第一个维度,大小为10,为batch维度。我的torch版本是1.7.1torch.autograd.grad没有参数is_grads_batched.

看来我是通过将函数cal_grad的返回码修改为:

解决了这个问题
    return torch.autograd.grad(loss_v, v, torch.ones_like(loss_v), create_graph=True, retain_graph=True)[0]