简单讲解pytorch使用GPU运行

一、前言

我们在pytorch深度学习进行模型预测的时候，往往会碰到数据庞大，模型复杂的情况，此时直接训练会调用CPU，速度很慢。我们可以选择使用GPU进行训练，提高效率。

二、方法

关于如何使用GPU网上目前没有比较简单的讲解，大多数都是一些例子或者复杂的情况。下面就根据自己的学习和实践来简单的讲一讲。

1、在训练数据和训练标签，测试数据和测试标签后面加上.cuda()

train_data = train_data.cuda()  
train_labels = train_labels.cuda()
test_data = test_data.cuda()
test_labels = test_labels.cuda()

2、将模型移动到cuda

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
lstm_model.to(device) # 移动模型到cuda

注意：其中lstm_model为你自己的模型实例化对象的名字。

3、在转化为numpy前要先将数据移动到cpu

predictive_y_for_training = predictive_y_for_training.view(-1, 1).data.cpu().numpy()

在前面.cuda()后,数据变为CUDA tensor.而numpy不能读取CUDA tensor，需要将它转化为 CPU tensor