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

Question

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

实施#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)

我想问的是：

第二个实现如我所料运行良好，但为什么第一个不训练 self.A（具有 None 梯度）？我猜 torch.autograd.grad 本身可能无法为其内部变量区分，因为外部可学习梯度下降步长 self.rho 和输入 x_init 获得了它们自己的 grads.
要实现更复杂的管道，其中self.A甚至可能是一个大网络，我该怎么做才能实现#1与 #2 一样有效（即，使 self.A 可以使用 autograd 机制进行训练）？
顺便说一句，我不确定x_init的grad是否在#1和[的转发中计算出来=59=]#2相同或完全相等。你能帮我检查一下吗？

注：x_init、y、x_gt的第一个维度，大小为10，为batch维度。我的torch版本是1.7.1，torch.autograd.grad没有参数is_grads_batched.

Answer 1

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

解决了这个问题

    return torch.autograd.grad(loss_v, v, torch.ones_like(loss_v), create_graph=True, retain_graph=True)[0]

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

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

python

deep-learning

pytorch

computer-vision

autograd