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

Posted 2023-03-12

【中文标题】pytorch 中的 autograd 可以处理同一模块中层的重复使用吗？

【英文标题】：Can autograd in pytorch handle a repeated use of a layer within the same module?

【发布时间】：2020-06-20 09:37:48

【问题描述】：

假设我在一个火炬模块中有一个层layer，并在单个forward 步骤中使用它两次或更多次，这样layer 输出的结果稍后再次输入到相同的layer。 pytorch的autograd能正确计算出这一层的权重grad吗？

这是我正在谈论的内容：

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 来处理这种情况，但比我的例子复杂得多吗？当一个层被迭代调用时，它是如何工作的？

【问题讨论】：

Autograd 会处理得很好。这是循环网络和权重共享网络的标准方法。正如数学所指示的那样，在这种情况下，反向传播将为 layer 的权重累积两次梯度。

【参考方案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 的本教程。