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

扎眼的阳光 pytorch 193

原文标题 ：Is this a correct way to training U-Net using different GPUs

U网代码：

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 脚本的最佳方法是什么？

原文链接：https://stackoverflow.com//questions/71936047/is-this-a-correct-way-to-training-u-net-using-different-gpus

回复

我来回复
  • Deusy94的头像
    Deusy94 评论

    您的代码应该可以工作，但我建议使用某种变量将子模型/张量传输到不同的 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)

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

    2年前 0条评论

相关问题