1. 查看是否有可用的GPU
use_gpu = torch.cuda.is_available()
use_gpu 为True则表示有GPU可用
2. 将网络的模型和loss函数放到GPU中
loss_fc = nn.MSELoss()
# While constructing the network, transfer the model to GPU (pytorch)
if (use_gpu):
model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
# model = model.cuda()
loss_fc = loss_fc.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate, betas=(0.9, 0.99), weight_decay=w_decay)
3. 将训练的数据数据放入GPU中,数据类型需为Tensor
for i, data in enumerate(data_loader):
train_x, train_y = data
# While traning, transfer the data to GPU
if (use_gpu):
train_x, train_y = train_x.cuda(), train_y.cuda()
4. 最后从GPU中取出需要最后输出的数据
if (use_gpu):
loss,prediction_y= loss.cpu(),prediction_y.cpu()