如何训练 PyTorch 中的共享层

Question

我有以下代码

import torch
import torch.nn as nn
from torchviz import make_dot, make_dot_from_trace

class Net(nn.Module):
    def __init__(self, input, output):
        super(Net, self).__init__()
        self.fc = nn.Linear(input, output)

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

model = Net(12, 12)
print(model)

x = torch.rand(1, 12)
y = model(x)
make_dot(y, params = dict(model.named_parameters()))

在这里，我在forward中重用了self.fc两次。

计算图是看

我对计算图感到困惑，我很好奇如何在反向传播中训练这个 model？ 在我看来，渐变将永远存在于一个循环中。 非常感谢。

Answer 1

您的图表没有问题。 您可以像任何其他前馈 model 一样训练它。

1. 关于循环：由于它是有向无环图，因此没有实际循环（查看箭头方向）。
2. 关于反向传播：让我们考虑fc.bias参数。 由于您重复使用同一层两次，因此偏差有两个向外的箭头（用于网络的两个地方）。 在反向传播阶段，方向是相反的：偏差会从两个地方得到梯度，这些梯度会加起来。
3. 关于图： FC 层可以表示为： Addmm(bias, x, T(weight) ，其中T是转置， Addmm是矩阵乘法加上向量。因此，您可以看到数据（ weight ， bias ）如何被传递给函数（ Addmm ， T ）

• https://pytorch.org/docs/stable/generated/torch.addmm.html
• https://pytorch.org/docs/stable/generated/torch.t.html