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

我有一层layer in an 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))

Output:

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 引擎的角度来看，这不是循环应用程序，因为生成的计算图会将重复计算展开为线性序列。为了说明这一点，对于单个层，您可能有：

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

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

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

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

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

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

循环神经网络在其基础上使用这种层（单元）的重复应用。例如，请参阅本教程使用字符级 RNN 对名称进行分类 https://pytorch.org/tutorials/intermediate/char_rnn_classification_tutorial.html.