修改现有 Pytorch 代码以在多个 GPU 上运行

Question

我正在尝试从以下链接在 2 个或更多 GPU 上运行 Pytoch UNet

Pytorch-UNet github

我到现在所做的改变是：

1.来自：

net = UNet(n_channels=3, n_classes=1, bilinear=True)
logging.info(f'Network:\n'
             f'\t{net.module.n_channels} input channels\n'
             f'\t{net.module.n_classes} output channels (classes)\n'
             f'\t{"Bilinear" if net.module.bilinear else "Transposed conv"} upscaling')

到：

net = UNet(n_channels=3, n_classes=1, bilinear=True)
net = nn.DataParallel(net)
logging.info(f'Network:\n'
             f'\t{net.module.n_channels} input channels\n'
             f'\t{net.module.n_classes} output channels (classes)\n'
             f'\t{"Bilinear" if net.module.bilinear else "Transposed conv"} upscaling')

在每个地方：

net.<something>

替换为：

net.module.<something>

我知道 pytorch 看到的 GPU 多于 1 个，因为torch.cuda.device_count()返回

2

.

但是只要我尝试运行需要比第一个 GPU 所拥有的更多内存的火车：

运行时错误：CUDA 内存不足。 尝试分配 512.00 MiB（GPU 0；11.91 GiB 总容量；10.51 GiB 已分配；82.56 MiB 空闲；818.92 MiB 缓存）

我通过改变批量大小来改变训练所需的内存。 欢迎任何帮助

编辑

我看到使用 2 个 GPU 训练运行速度快两倍，但使用单个 GPU 运行的最大批量大小与两个 GPU 相同。 有没有什么办法可以在一次训练中同时使用 2 个 GPU 的内存？

Answer 1

我的错误是将output = net(input) （通常命名为model ）更改为：

output = net.module(input)

你可以在这里找到信息