在多个 GPU 上训练单个 pytorch 模型并修复某些层？

Question

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

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

   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. 但是有两个训练阶段。在第一阶段，我想修复layer0和layer1的B。在第二阶段，只修复layer0。那么layer1中的参数requires_grad应该在训练期间改变。但是，DistributedDataParalleldoc 说：

   您应该切勿在使用 DistributedDataParallel 封装模型后尝试更改模型的参数。

实际上，我尝试使用B.module 来引用包裹在B 中的A。但与单GPU模型相比，测试结果异常。也许这种方式是不允许的。

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

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

2019.12.03 更新

按照@jodag 的建议，我尝试了DataParallel，但没有成功。这次我在包裹之后没有改变 B 中的任何东西（除了训练它）。为简单起见，我的代码是这样的（我引用了this）：

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 中）并仅将其部分参数传递给相关的优化器。