指定GPU训练

终端命令

# 等号前后不允许存在空格
CUDA_VISIBLE_DEVICES=0,1,2 python train.py

代码中指定

# 放在import torch前面,在一处指定即可,多处的话需要保持一致(比如多处import os)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"

其他

import torch
 
torch.cuda.is_available()  # 查看cuda是否可用
 
torch.cuda.device_count()  # 返回GPU数目
 
torch.cuda.get_device_name(0)  # 返回GPU名称,设备索引默认从0开始
 
torch.cuda.current_device()  # 返回当前设备索引
最后修改:2022 年 06 月 26 日
如果觉得我的文章对你有用,请随意赞赏