一、数据、模型、损失函数需要迁移到GPU上

使用GPU训练时，数据、函数和模型都必须同时放在GPU上，否则会出错。

#（1）判断GPU是否可用
if torch.cuda.is_available():
    device = torch.device('cuda')
    cudnn.benchmark = True
else:
    device = torch.device('cpu')
 
#（2）构建模型时，把网络，与损失函数迁移到GPU上
model = CNN().to(device)
loss = nn.CrossEntropyLoss().to(device)
 
#（3）训练模型时，把数据迁移到GPU上
x, y = x.to(device), y.to(device)
 
# 注，只有tensor类型才能上传到GPU上，故需要对numpy数据进行tensor类型转换
# torch.tensor(x) 或 torch.from_numpy(x) 

 
#（4）对训练的输出结果有些需要使用np的函数进行操作
# 故需先将输出结果迁移CPU上，再转成numpy类型，然后使用np的函数
output = (model(x)).cpu().numpy()

二、知识拓展

2.1 指定单个GPU

import os
# 按照PCI_BUS_ID顺序从0开始排列GPU设备 
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
model.cuda()

2.2 指定多个GPU

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

device_ids = [0,1,2,3]
# 默认使用所有的device_ids
model  = torch.nn.Dataparallel(model, device_ids =device_ids)
model.cuda()

2.3 .to(device)方法

# 单GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 多GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.DataParallel(model，device_ids=[0,1,2])
model.to(device)

2.4 将模型放入GPU的两种方式

# 方式1
model1 = Model(para1=1,para2=2)
model1.cuda()

# 方式2（推荐，更简洁）
model2 = Model(para1=1,para2=2).cuda()

注意：在将模型实例放到GPU上时，即用model.cuda()时，只会将__init__中的有self前缀的属性及函数放到GPU上，对于其他的函数无作用。所以非_init__中定义的函数和数据需要单独的放入GPU语句，也就是上述代码中的para1和para2参数。