pytorch 使用gpu

1. pytorch查看设备信息

# 查看gpu设备是否可用
torch.cuda.is_available() 

# 返回gpu数量；
torch.cuda.device_count()

# 返回gpu名字，设备索引默认从0开始；
torch.cuda.get_device_name(0)

# 返回当前设备索引；
torch.cuda.current_device()

2. 设置device

一般我们设置第一个gpu（主要是因为我只有一块gpu。。。）

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

然后我们就可以使用to方法快捷的将数据放在gpu中计算：

model = CNN().to(device) # model放入gpu
criteon = nn.CrossEntropyLoss().to(device) # 损失函数
data, target = data.to(device), target.cuda() # 训练的数据

这里建议使用 to(device)，而不是 cuda()，毕竟 device 全局设置可以方便管理，且对于没有gpu的设备，device可以设置为cpu。