PyTorch torch.cuda.device_count 返回值与实际 GPU 数量不一致

今天执行基于 PyTorch 的图像分类算法程序时,触发了自己写的断言错误。而断言的细节,就是判断用户输入的 GPU 编号是否合法。
调试打开,发现 torch.cuda.device_count() 返回的是 1。而我机器上明明是两张卡。

一脸懵逼。
查阅 PyTorch 官网后,发现是使用问题。我在调用 device_count 之前,已经设置过了环境变量 CUDA_VISIBLE_DEVICES
通过在 os.environ["CUDA_VISIBLE_DEVICES"] 代码之前执行 device_count, 发现返回的是 2。至此,问题已定位。
PS. 官方推进使用 os.environ["CUDA_VISIBLE_DEVICES"] 的形式来设定使用的 GPU 显卡。

参考链接

posted @ 2020-04-20 19:15  Scharfsinnig  阅读(23434)  评论(0编辑  收藏  举报