如何在 Pytorch 中将神经网络划分为子网络？

Question

我想使用 Pytorch 将一个神经网络划分为两个子网络。为了使事情具体化，请考虑以下图像：

在 1 中，我有一个 3x4x1 的神经网络。我想要的是，例如在 epoch 1 期间，我只想更新子网络 1 中的权重，即必须冻结子网络 2 中出现的权重。再说一次，在 epoch 2 中，我想训练出现在子网络 2 中的权重，而其余的应该被冻结。

我该怎么做？

Answer 1

如果您的子网是层的子集，您可以轻松做到这一点。也就是说，您不需要冻结任何部分图层。要么全有，要么全无。

对于您的示例，这意味着将隐藏层分为两个不同的 2 节点层。每一个都属于一个子网，这让我们回到全部或一无所有。

完成后，您可以使用requires_grad 切换各个层。在参数上将此设置为False 将禁用训练并冻结权重。要对整个模型、子模型或 Module 执行此操作，请循环访问 model.parameters()。

对于您的示例，有 3 个输入、1 个输出和一个现在拆分的 2x2 隐藏层，它可能看起来像这样：

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

def set_grad(model, grad):
    for param in model.parameters():
        param.requires_grad = grad

class HalfFrozenModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.hid1 = torch.nn.Linear(3, 2)
        self.hid2 = torch.nn.Linear(3, 2)
        self.out = torch.nn.Linear(4, 1)

    def set_freeze(self, hid1=False, hid2=False):
        set_grad(self.hid1, not hid1)
        set_grad(self.hid2, not hid2)

    def forward(self, inp):
        hid1 = self.hid1(inp)
        hid2 = self.hid2(inp)
        hidden = torch.cat([hid1, hid2], 1)
        return self.out(F.relu(hidden))

然后你可以像这样训练一个或另一个：

model = HalfFrozenModel()
model.set_freeze(hid1=True)
# Do some training.
model.set_freeze(hid2=True)
# Do some more training.
# ...

---

如果您碰巧使用fastai，那么还有一个图层组的概念也用于此目的。 fastai documentation 详细介绍了它的工作原理。