这是使用不同gpu训练U-Net的正确方法吗?

Question

U-Net代码:

class UNet(nn.Module):
def __init__(self):
super(UNet, self).__init__()
self.c1 = convBlock(1, 64).to('cuda:0')
self.d1 = downSample(64).to('cuda:0')
self.c2 = convBlock(64, 128).to('cuda:0')
self.d2 = downSample(128).to('cuda:0')
self.c3 = convBlock(128, 256).to('cuda:0')
self.d3 = downSample(256).to('cuda:1')
self.c4 = convBlock(256, 512).to('cuda:1')
self.d4 = downSample(512).to('cuda:1')
self.c5 = convBlock(512, 1024).to('cuda:1')
self.u1 = upSample(1024).to('cuda:1')
self.c6 = convBlock(1024, 512).to('cuda:1')
self.u2 = upSample(512).to('cuda:1')
self.c7 = convBlock(512, 256).to('cuda:1')
self.u3 = upSample(256).to('cuda:1')
self.c8 = convBlock(256, 128).to('cuda:1')
self.u4 = upSample(128).to('cuda:0')
self.c9 = convBlock(128, 64).to('cuda:0')
self.out = nn.Conv3d(64, 1, 3, 1, 1).to('cuda:0')
self.th = nn.Sigmoid().to('cuda:0')
def forward(self, x):
L1 = self.c1(x.to('cuda:0'))
L2 = self.c2(self.d1(L1.to('cuda:0')).to('cuda:0'))
L3 = self.c3(self.d2(L2.to('cuda:0')).to('cuda:0'))
L4 = self.c4(self.d3(L3.to('cuda:1')).to('cuda:1'))
L5 = self.c5(self.d4(L4.to('cuda:1')).to('cuda:1'))
R4 = self.c6(self.u1(L5.to('cuda:1'), L4.to('cuda:1')).to('cuda:1'))
R3 = self.c7(self.u2(R4.to('cuda:1'), L3.to('cuda:1')).to('cuda:1'))
R2 = self.c8(self.u3(R3.to('cuda:1'), L2.to('cuda:1')).to('cuda:1'))
R1 = self.c9(self.u4(R2.to('cuda:0'), L1.to('cuda:0')).to('cuda:0'))
return self.th(self.out(R1.to('cuda:0')).to('cuda:0'))

convBlock, downSample, upSample是我自己代码中的图层。

我想训练3DU-Net，但是GPU内存不够，所以我想用多个GPU来训练这个模型。

我将不同的U-net层分配给不同的gpu。

我想问，如果这是正确的方式使用不同的gpu来训练模型?在Linux服务器中使用PyTorch模块运行多个GPU训练python脚本的最佳方法是什么?

Answer 1

你的代码应该工作，但我建议使用某种变量来转移子模型/张量到不同的gpu。我一直用的是这样的:

class MyModel(nn.Module):
def __init__(self, split_bool: bool = False):
self.submodule1 = ...
self.submodule2 = ...
self.split_bool = split_bool
if split_bool:
self.submodule1.cuda(0)
self.submodule2.cuda(1)
def forward(self, x):
x = self.submodule1(x)
if self.split_bool:
x = x.cuda(1) # Transfer tensor to second GPU
return self.submodule2(x)

对于多次训练，这真的取决于你的服务器。您是否使用tensorboard/tensorboardX来绘制结果?您可以使用tmux启动具有不同参数的多个训练脚本，甚至可以编写自己的bash脚本。