python - pytorch中的autograd可以处理同一模块中层的重复使用吗？

Question

我有一层 layer在 nn.Module并在一次 forward 中使用两次或更多次步。这个 layer 的输出稍后输入到相同的layer .可以pytorch的autograd正确计算该层权重的梯度？

def forward(x):
    x = self.layer(x)
    x = self.layer(x)
    return x

完整示例:

import torch
import torch.nn as nn
import torch.nn.functional as F

class net(nn.Module):
    def __init__(self,in_dim,out_dim):
        super(net,self).__init__()
        self.layer = nn.Linear(in_dim,out_dim,bias=False)

    def forward(self,x):
        x = self.layer(x)
        x = self.layer(x)
        return x

input_x = torch.tensor([10.])
label = torch.tensor([5.])
n = net(1,1)
loss_fn = nn.MSELoss()

out = n(input_x)
loss = loss_fn(out,label)
n.zero_grad()
loss.backward()

for param in n.parameters():
    w = param.item()
    g = param.grad

print('Input = %.4f; label = %.4f'%(input_x,label))
print('Weight = %.4f; output = %.4f'%(w,out))
print('Gradient w.r.t. the weight is %.4f'%(g))
print('And it should be %.4f'%(4*(w**2*input_x-label)*w*input_x))

输出:

Input = 10.0000; label = 5.0000
Weight = 0.9472; output = 8.9717
Gradient w.r.t. the weight is 150.4767
And it should be 150.4766

在这个例子中，我定义了一个只有一个线性层( in_dim=out_dim=1 并且没有偏差)的模块。 w是该层的权重； input_x是输入值； label是期望的值。由于损失选择为 MSE，因此损失的公式为 ((w^2)*input_x-label)^2手工计算，我们有 dw/dx = 2*((w^2)*input_x-label)*(2*w*input_x)我上面示例的输出显示 autograd给出与手动计算相同的结果，让我有理由相信它可以在这种情况下工作。但在实际应用中，该层可能具有更高维度的输入和输出，之后是非线性激活函数，并且神经网络可以具有多个层。
我想问的是:我可以信任 autograd处理这种情况，但比我的例子复杂得多？当一个层被迭代调用时它是如何工作的？

Answer 1

这将工作得很好。从 autograd 引擎的角度来看，这不是一个循环应用程序，因为生成的计算图会将重复计算展开为线性序列。为了说明这一点，对于单层，您可能有:

x -----> layer --------+
           ^           |
           |  2 times  |
           +-----------+

从 autograd 的角度来看，这看起来像:

x ---> layer ---> layer ---> layer

这里 layer是在图上复制 3 次的同一层。这意味着在计算层权重的梯度时，它们将从所有三个阶段累积。所以当使用 backward :

x ---> layer ---> layer ---> layer ---> loss_func
                                            |
       lback <--- lback <--- lback <--------+
         |          |          |
         |          v          |
         +------> weights <----+
                   _grad

这里 lback表示 layer 的局部导数使用上游梯度作为输入的正向变换。每一个都添加到图层的 weights_grad .

循环神经网络在其基础上使用层(单元)的这种重复应用。例如，参见本教程关于 Classifying Names with a Character-Level RNN .