pytorch - 在pytorch中构造参数组

Question

在 torch.optim 中 documentation ，据说可以使用不同的优化超参数对模型参数进行分组和优化。据说

For example, this is very useful when one wants to specify per-layerlearning rates:

optim.SGD([
                {'params': model.base.parameters()},
                {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

This means that model.base’s parameters will use the defaultlearning rate of 1e-2, model.classifier’s parameters will use alearning rate of 1e-3, and a momentum of 0.9 will be used for allparameters.

我想知道如何定义具有 parameters() 属性的组。我想到的是以下形式的东西

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.base()
        self.classifier()

        self.relu = nn.ReLU()

    def base(self):
        self.fc1 = nn.Linear(1, 512)
        self.fc2 = nn.Linear(512, 264)

    def classifier(self):
        self.fc3 = nn.Linear(264, 128)
        self.fc4 = nn.Linear(128, 964)

    def forward(self, y0):

        y1 = self.relu(self.fc1(y0))
        y2 = self.relu(self.fc2(y1))
        y3 = self.relu(self.fc3(y2))

        return self.fc4(y3)

我应该如何修改上面的代码片段才能获得 model.base.parameters()？是定义 nn.ParameterList 并将所需层的 weight 和 bias 显式添加到该列表的唯一方法吗？最佳做法是什么？

Answer 1

我将展示三种解决此问题的方法。不过最终还是要看个人喜好。

---

- 使用 nn.ModuleDict 对参数进行分组.

我注意到这里有一个答案使用 nn.Sequential 对允许的层进行分组使用 nn.Sequential 的 parameters 属性定位模型的不同部分。事实上，base 和分类器可能不仅仅是顺序层。我相信更通用的方法是让模块保持原样，而是初始化一个额外的 nn.ModuleDict 模块，它将包含优化组在单独的 nn.ModuleList 中排序的所有参数。小号:

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.fc1 = nn.Linear(1, 512)
        self.fc2 = nn.Linear(512, 264)
        self.fc3 = nn.Linear(264, 128)
        self.fc4 = nn.Linear(128, 964)

        self.params = nn.ModuleDict({
            'base': nn.ModuleList([self.fc1, self.fc2]),
            'classifier': nn.ModuleList([self.fc3, self.fc4])})

    def forward(self, y0):
        y1 = self.relu(self.fc1(y0))
        y2 = self.relu(self.fc2(y1))
        y3 = self.relu(self.fc3(y2))
        return self.fc4(y3)

然后你可以定义你的优化器:

optim.SGD([
    {'params': model.params.base.parameters()},
    {'params': model.params.classifier.parameters(), 'lr': 1e-3}
], lr=1e-2, momentum=0.9)

请注意 MyModel 的 parameters 生成器不会包含重复参数。

---

- 创建用于访问参数组的接口(interface)。

一个不同的解决方案是在 nn.Module 中提供一个接口(interface)来将参数分成几组:

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.fc1 = nn.Linear(1, 512)
        self.fc2 = nn.Linear(512, 264)
        self.fc3 = nn.Linear(264, 128)
        self.fc4 = nn.Linear(128, 964)

    def forward(self, y0):
        y1 = self.relu(self.fc1(y0))
        y2 = self.relu(self.fc2(y1))
        y3 = self.relu(self.fc3(y2))
        return self.fc4(y3)

    def base_params(self):
        return chain(m.parameters() for m in [self.fc1, self.fc2])

    def classifier_params(self):
        return chain(m.parameters() for m in [self.fc3, self.fc4])

已导入 itertools.chain作为链。

然后定义你的优化器:

optim.SGD([
    {'params': model.base_params()},
    {'params': model.classifier_params(), 'lr': 1e-3}
], lr=1e-2, momentum=0.9)

---

- 使用子 nn.Module。

最后，您可以将模块部分定义为子模块(这里归结为 nn.Sequential 方法，但您可以将其推广到任何子模块)。

class Base(nn.Sequential):
    def __init__(self):
        super().__init__(nn.Linear(1, 512),
                         nn.ReLU(),
                         nn.Linear(512, 264),
                         nn.ReLU())

class Classifier(nn.Sequential):
    def __init__(self):
        super().__init__(nn.Linear(264, 128),
                         nn.ReLU(),
                         nn.Linear(128, 964))

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.base = Base()
        self.classifier = Classifier()

    def forward(self, y0):
        features = self.base(y0)
        out = self.classifier(features)
        return out

在这里您可以再次使用与第一种方法相同的界面:

optim.SGD([
    {'params': model.base.parameters()},
    {'params': model.classifier.parameters(), 'lr': 1e-3}
], lr=1e-2, momentum=0.9)

我认为这是最佳实践。但是，它迫使您将每个组件定义到单独的 nn.Module 中，这在试验更复杂的模型时可能会很麻烦。