pytorch中有关gpu的操作

与GPU的相关操作

1. 查看gpu是否可用

torch.cuda.is_available()      # 可用返回True，不可用返回False

2. 返回可用的CPU的数量

torch.cuda.device_count()

3. 返回gpu的名字，设备索引默认从0开始

torch.cuda.get_device_name(0)

4. 返回当前设备索引

torch.cuda.current_device()

5. 使用gpu进行训练计算

# 有两种方法
# 第一种方法：
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)   # 将模型移到gpu上去
data = data.to(device)     # 将数据移到gpu上去

# 第二种方法：
model= model.cuda()
data = data.cuda()