python - 在多个 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. 但是有两个训练阶段。在第一阶段，我想修复B的layer0和layer1。第二阶段，仅修复layer0。然后在训练过程中应该更改layer1中参数的requires_grad。但是，DistributedDataParallel doc说:

   You should never try to change your model’s parameters after wrapping up your model with 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 中)并仅将其部分参数传递给相关优化器。