这是使用不同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来训练这个模型。

我给不同的GPU分配不同的U-net层.

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

Answer 1

您的代码应该可以工作，但我建议使用某种变量将子模型/张量传输到不同的gpus。这就是我一直在用的东西：

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)

对于多个培训，它实际上取决于您的服务器。您是否使用张卡/张卡来绘制结果？您可以使用tmux启动具有不同参数的多个培训脚本，甚至可以编写自己的bash脚本。