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

Question

优网代码：

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

您的代码应该可以工作，但我建议使用某种变量将 submodel/tensor 传输到不同的 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 脚本。

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

Is this a correct way to training U-Net using different GPUs

gpu

pytorch