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

Question

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

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

我怎样才能做到这一点？

Answer 1

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

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

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

对于您的示例，有 3 个输入、1 个 output 和一个现在拆分的 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 文档详细介绍了它的工作原理。