在 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

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 并重新安装它，这解决了问题。

Training a model on GPU is very slow