时间戳

首页 > 解决方案 > torch.cuda.is_available() 返回 true 但 Torch 模型继续在 CPU 上进行训练

deep-learning - torch.cuda.is_available() 返回 true 但 Torch 模型继续在 CPU 上进行训练

问题描述

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

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

标签： deep-learninggputorch

解决方案

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

推荐阅读