PyTorch指定GPU

 

推荐通过os环境设置可见的GPU编号。需要注意的是import  torch应该在os.environ后面。之前没有注意，发现一直用的是GPU 0。

import os
 
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3"
 
import torch

　　

查看GPU信息的一些常用指令

import torch
 
torch.cuda.is_available()  # 查看cuda是否可用
 
torch.cuda.device_count()  # 返回GPU数目
 
torch.cuda.get_device_name(0)  # 返回GPU名称，设备索引默认从0开始
 
torch.cuda.current_device()  # 返回当前设备索引