在GPU上训练模型非常缓慢

Question

我正在使用A100-SXM4-40GB Gpu，但训练非常缓慢。我尝试了两种模式，一种是cifar的简单分类，另一种是城市景观的Unet。我在其他GPU上尝试了我的代码，它运行得很好，但我不知道为什么在这个高容量GPU上训练非常慢。

如果有任何帮助，我将不胜感激。

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

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

• 英伟达驱动程序版本：460.32.03
• cuDNN版本：无法收集

Answer 1

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

Answer 2

初始化期间在模型上调用.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')