背景

实验室服务器有多张显卡，但今天“0”卡显存已被占满，因此我在代码中添加os.environ['CUDA_VISIBLE_DEVICES']=“2”指定使用编号为"2"的显卡(显存够的前提下)。但跑VGG16的时候却报错：

经过调试发现os.environ['CUDA_VISIBLE_DEVICES']想要生效，import torch的位置需要放在指令os.environ['CUDA_VISIBLE_DEVICES']后面才行。

验证方案

利用torch.cuda.device_count()协助判断指令os.environ['CUDA_VISIBLE_DEVICES']="x" 是否生效

# 实验室服务器有8台GPU
# 选择调用其中三台
# torch.cuda.device_count(): 查看可用的GPU数量

# 实验1. 将import torch放在指令os.environ['CUDA_VISIBLE_DEVICES']="x"前面
import os
import torch
os.environ['CUDA_VISIBLE_DEVICES']='0, 1, 2' # 程序可见的GPU
print(torch.cuda.device_count())  # 8  说明指令没生效

# 实验2. 将import torch放在指令os.environ['CUDA_VISIBLE_DEVICES']="x"后面
import os
os.environ['CUDA_VISIBLE_DEVICES']='0, 1, 2' # 程序可见的GPU
import torch
print(torch.cuda.device_count())  # 3  说明指令生效

总结

上述验证实验表明，如果想要os.environ['CUDA_VISIBLE_DEVICES']起作用，需将import torch放在指令os.environ['CUDA_VISIBLE_DEVICES']="x"后面