使用pytorch DistributedDataParallel时遇到一些问题。情况是:


我的模型是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

现在我有了一些新数据。我想在多个GPU上对其进行微调。我需要将A包装为多GPU模型B。
但是有两个训练阶段。在第一阶段中,我要修复B的layer0layer1。在第二阶段中,仅要修复layer0。然后在训练期间应更改requires_grad中参数的layer1。但是,DistributedDataParallel doc说:


  在用DistributedDataParallel包装模型之后,您永远不要尝试更改模型的参数。



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

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

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

更新2019.12.03

正如@jodag所建议的那样,我尝试了DataParallel,但没有成功。这次,我在包装B之后没有对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()


最佳答案

如果仅尝试优化部分参数,为什么不尝试通过优化器而不是模型来控制它呢?
您可以按原样保留模型(包装在DistributedDataParallel中),仅将其部分参数传递给相关的优化器。

关于python - 在多个固定层的GPU上训练单个pytorch模型?,我们在Stack Overflow上找到一个类似的问题:https://stackoverflow.com/questions/59134785/

10-11 22:48
查看更多