在 GPU 上训练模型非常慢

Question

我正在使用 A100-SXM4-40GB Gpu，但训练速度非常慢。我尝试了两种模型，一种是在 cifar 上的简单分类，另一种是在 Cityscapes 上的 Unet。我在其他 GPU 上尝试了我的代码，它运行得非常好，但我不知道为什么在这个大容量 GPU 上训练速度超级慢。

我将不胜感激。

以下是 GPU 的一些其他属性。

GPU 0: A100-SXM4-40GB
GPU 1: A100-SXM4-40GB
GPU 2: A100-SXM4-40GB
GPU 3: A100-SXM4-40GB

• Nvidia 驱动程序版本：460.32.03
• cuDNN 版本：无法收集

Answer 1

在初始化期间在模型上调用 .cuda()。

根据您的上述评论，您已经安装了 GPU 和 CUDA，因此没有必要使用 torch.cuda.is_available() 检查设备可用性。

此外，您应该将模型包装在 nn.DataParallel 中，以允许 PyTorch 使用您向其公开的每个 GPU。您也可以做 DistributedDataParallel，但 DataParallel 最初更容易掌握。

示例初始化：

model = UNet().cuda()
model = torch.nn.DataParallel(model)

此外，您可以通过执行带有以下标志的 python 脚本来确保将代码暴露给所有 GPU：

CUDA_VISIBLE_DEVICES=0,1,2,3 python3 train_unet.py

最后要注意 - nn.DataParallel 封装了模型本身，因此为了保存 state_dict，您需要访问 DataParallel 中的模块：

torch.save(model.module.state_dict(), 'unet.pth')

Answer 2

感谢您的回答。在尝试您的答案之前，我决定卸载 anaconda 并重新安装它，这解决了问题。