在pytorch中构建参数组

Question

在torch.optimdocumentation中说明模型参数可以分组，用不同的优化超参数进行优化。它说

For example, this is very useful when one wants to specify per-layer learning 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 default learning rate of 1e-2, model.classifier’s parameters will use a learning rate of 1e-3, and a momentum of 0.9 will be used for all parameters.

我想知道如何定义具有 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

您可以使用torch.nn.Sequential来定义base和classifier。您的 class 定义可以是：

class MyModel(nn.Module):

    def __init__(self):
        super(MyModel, self).__init__()
        self.base = nn.Sequential(nn.Linear(1, 512), nn.ReLU(), nn.Linear(512,264), nn.ReLU())
        self.classifier = nn.Sequential(nn.Linear(264,128), nn.ReLU(), nn.Linear(128,964))

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

然后，您可以使用 model.base.parameters() 和 model.classifier.parameters() 访问参数。

Answer 2

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

- 使用 `nn.ModuleDict`.

对参数进行分组

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

中排序的所有参数

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's parameters' 生成器不会包含重复参数。

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

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

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 导入为 chain。

然后定义你的优化器：

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

- 使用 child `nn.Module`s.

最后，您可以将模块部分定义为子模块（这里归结为 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，这在试验更复杂的模型时可能会很麻烦。

在pytorch中构建参数组

Constructing parameter groups in pytorch

pytorch

- 使用 nn.ModuleDict.

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

- 使用 child nn.Modules.

- 使用 `nn.ModuleDict`.

- 使用 child `nn.Module`s.