今天执行基于 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 显卡。