PyTorch查看和指定GPU
文章目录1.查看GPU1.1.查看当前GPU是否可用1.2.查看GPU数量1.3.返回GPU名字1.4.返回当前设备索引2.指定GPU2.1.使用`CUDA_VISIBLE_DEVICES`2.1.1.在终端中指定2.1.2.代码中指定2.2.使用函数`set_device`1.查看GPU1.1.查看当前GPU是否可用#cuda是否可用torch.cuda.is_available()# 返回 T
·
文章目录
1.查看GPU
1.1.查看当前GPU是否可用
#cuda是否可用
torch.cuda.is_available() # 返回 True 则可用, 返回 False 则不可用
1.2.查看GPU数量
# 返回gpu数量
torch.cuda.device_count() # 返回值就是 GPU 数目
1.3.返回GPU名字
# 返回gpu名字,设备索引默认从0开始
torch.cuda.get_device_name(index) # index 是索引, 默认从 0 开始
1.4.返回当前设备索引
# 返回当前设备索引
torch.cuda.current_device()
2.指定GPU
2.1.使用 CUDA_VISIBLE_DEVICES
2.1.1.在终端中指定
CUDA_VISIBLE_DEVICES=1 python main.py
2.1.2.代码中指定
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
上述两种方式其实都是添加环境变量,限制CUDA程序所能使用的GPU设备。从名字上,就能看出是限定程序能“看到”的设备。
CUDA应用运行时,CUDA将遍历当前可见的设备,并从零开始为可见设备编号。
可以同时指定多个设备(设备是真实存在的,不要指定不存在的设备)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3"
2.2.使用函数 set_device
import torch
torch.cuda.set_device(id)
官方建议使用CUDA_VISIBLE_DEVICES
,不建议使用 set_device
函数。
更多推荐
已为社区贡献3条内容
所有评论(0)