torch.cuda.is_available() 返回 true，但火炬模型继续在 CPU 上进行训练 [关闭]

Posted 2023-03-29

【中文标题】torch.cuda.is_available() 返回 true，但火炬模型继续在 CPU 上进行训练 [关闭]

【英文标题】：torch.cuda.is_available() returns true but torch models keep training on CPU [closed]

【发布时间】：2021-06-21 22:35:35

【问题描述】：

我尝试创建一个新环境，如下所示： https://github.com/microsoft/computervision-recipes/blob/master/SETUP.md

我检查了 cuda 运行良好，检测到我的 gpu 并且一切看起来都很好。但是当我拟合模型时，nvdia-smi 显示没有占用 GPU 并且 CPU 为 100%。

【参考方案1】：

Torch 需要为每个模型和训练过程设置 Cuda 设备，您需要手动将所有权重和另一个张量传递给 Cuda 设备。像这样：

a = torch.tensor([1, 2, 3])
a = a.to('cuda')

但在您的情况下，也许您需要使用库中的方法。我建议检查这个https://github.com/microsoft/computervision-recipes/blob/1bb489af757fde7c773e16fab87b24305cff4457/utils_cv/tracking/references/fairmot/trains/base_trainer.py#L32