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

Question

我在 nn.Module 中有一层 layer 并在单个 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.

的本教程

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

Can autograd in pytorch handle a repeated use of a layer within the same module?

python

neural-network

pytorch

autograd