使用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的
layer0
和layer1
。在第二阶段中,仅要修复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/