在固定某些层的多个 GPU 上训练单个 pytorch 模型？

Question

我在使用 pytorch DistributedDataParallel时遇到了一些问题。 情况是：

1. 我的模型是A ，它像往常一样在单个GPU 上训练。 假设有三层答：

    class A(nn.module): def __init__(self): super(A,self).__init__() self.layer0 = layer0 self.layer1 = layer1 self.layer2 = layer2 def forward(self,x): x=self.layer0(x) x=self.layer1(x) x=self.layer2(x) return x

2. 现在我有了一些新数据。 我想在多个GPU 上用它微调A。 我需要将A包装为多 GPU 模型B 。

3. 但是有两个训练阶段。 在第一阶段，我想修复B 的layer0和layer1 。 在第二阶段，只修复layer0 。 然后requires_grad在训练期间更改layer1 requires_grad参数。 但是， DistributedDataParallel文档说：

   在使用 DistributedDataParallel 包装模型后，您永远不应该尝试更改模型的参数。

事实上，我试图用B.module指一个包裹在B中。 但与单GPU模型相比，测试结果异常。 也许这种方式是不允许的。

我该怎么办？ 有什么合适的方法来包装我的模型吗？ 保存和加载模型时应该注意什么？

只需在具有多个 GPU 的单台机器上运行它，这样您就可以忽略使用多台机器的分布式情况。 非常感谢。

2019.12.03 更新

正如@jodag 所建议的，我尝试了DataParallel ，但没有用。 这次我在包装之后没有改变B中的任何东西（除了训练它）。 为简单起见，我的代码是这样的（我参考了这个）：

class B(nn.DataParallel):
     def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.module, name)
a = A()
b = B(a,device_ids=[0,1])
b = b.cuda()
trained_param = b.layer2.parameters()
# trained_param = [{'params':b.layer2.parameters()},{'params':b.layer1.parameters()}]
optimizer = optim.Adam(trained_param)
b.train()
...
for x, label in data_loader:
    optimizer.zero_grad()
    x = x.to(0) # This line can be commented.
    y = b(x)
    l = loss(y, label)
    l.backword()
    optimizer.step()

Answer 1

如果您只尝试优化部分参数，为什么不尝试通过优化器而不是模型来控制它？
您可以保持模型原样（包装在DistributedDataParallel ）并仅将其部分参数传递给相关优化器。