一、前言

我们在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

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐