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

Posted 2023-03-28

【中文标题】修改现有 Pytorch 代码以在多个 GPU 上运行

【英文标题】：Modify existing Pytorch code to run on multiple GPUs

【发布时间】：2021-01-18 23:00:05

【问题描述】：

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

Pytorch-UNet github

到目前为止我所做的更改是：

1。 来自：

net = UNet(n_channels=3, n_classes=1, bilinear=True)
logging.info(f'Network:\n'
             f'\tnet.module.n_channels input channels\n'
             f'\tnet.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'\tnet.module.n_channels input channels\n'
             f'\tnet.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 需要更多内存的火车：

RuntimeError: CUDA 内存不足。尝试分配 512.00 MiB (GPU 0; 11.91 GiB 总容量；已分配 10.51 GiB； 82.56 MiB 免费； 818.92 MiB 缓存）

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

编辑

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

【问题讨论】：

降低批量大小，直到它适合最小数量的 GPU 内存。

@Rika，我想在 2 个 GPU 上运行，小批量运行没有问题

您的 GPU 必须具有相同数量的内存，如果它们的数量不同，那么 Pytorch 将使用较小的数量作为两个 GPU 上的可用 vram 数量。所以假设你有一个 8Gig 和一个 12Gig gpu。当你想训练时，你的批量大小应该不超过 8Gig。如果是这样，您将无法训练，并且会失败并显示您收到的错误消息

谢谢@Rika，我看到了使用 1 个和 2 个 GPU 运行之间的区别，我的意思是使用 2 个 GPU 运行速度要快两倍。但我不明白为什么它不使用第二个 GPU 的内存。我的意思是我应该怎么做才能使用第二个的内存，根据你的例子甚至是 16GB。

【参考方案1】：

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

output = net.module(input)

你可以找到信息here