pytorch学习之实现代码在GPU上运行

当模型太大，或者参数太多的情况下，为了加快训练速度，经常会使用GPU来进行训练

此时我们的代码需要稍作调整：

在cup上运行的代码：

#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import torch
from torch import nn
from torch import optim
import numpy as np
from matplotlib import pyplot as plt

#1、定义数据
x = torch.rand([50,1])
y = x*3 + 0.8

#定义模型
class Lr(nn.Module): #继承nn.Module
    def __init__(self):
        super(Lr, self).__init__()
        self.linear = nn.Linear(1,1)

    def forward(self,x):
        out = self.linear(x)
        return out
#2、实例化模型、loss函数以及优化器
model = Lr()
criterion = nn.MSELoss()   #损失函数
optimizer = optim.SGD(model.parameters(),lr=0.01) #优化器

#3、训练模型
for i in range(5000):
    out = model(x)# 获取预测值
    loss = criterion(y,out) #计算损失
    optimizer.zero_grad() #梯度归零
    loss.backward() #计算梯度
    optimizer.step() #更新梯度

    if(i+1) % 20 ==0:
        print('Epoch[{}/{}],loss:{:.6f}'.format(i,5000,loss.data))


#4、模型评估
model.eval() #设置模型为评估模式，即预测模式
predict = model(x)
predict = predict.data.numpy()
plt.scatter(x.data.numpy(),y.data.numpy(),c="r")
plt.plot(x.data.numpy(),predict)
plt.show()

1、判断GPU是否可用torch.cuda.is_available()

torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device(type='cuda', index=0)  #使用gpu
device(type='cpu') #使用cpu

2、把模型参数和input数据转化为cuda的支持类型

x = torch.rand([50,1]).to(device)
model = Lr().to(device)

3、在GPU上计算结果也为cuda的数据类型，需要转化为numpy或者torch的cpu的tensor类型

predict = predict.cpu().detach().numpy()
plt.scatter(x.cpu().detach().numpy(),y.cpu().detach().numpy(),c="r")
plt.plot(x.cpu().detach().numpy(),predict)

detach()的效果和data的相似，但是detach()是深拷贝，data是取值，是浅拷贝

修改后的代码：

#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import torch
from torch import nn
from torch import optim
import numpy as np
from matplotlib import pyplot as plt

#定义一个device对象
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#1、定义数据
x = torch.rand([50,1]).to(device)
y = x*3 + 0.8

#定义模型
class Lr(nn.Module): #继承nn.Module
    def __init__(self):
        super(Lr, self).__init__()
        self.linear = nn.Linear(1,1)

    def forward(self,x):
        out = self.linear(x)
        return out
#2、实例化模型、loss函数以及优化器
model = Lr().to(device)
criterion = nn.MSELoss()   #损失函数
optimizer = optim.SGD(model.parameters(),lr=0.01) #优化器

#3、训练模型
for i in range(5000):
    out = model(x)# 获取预测值
    loss = criterion(y,out) #计算损失
    optimizer.zero_grad() #梯度归零
    loss.backward() #计算梯度
    optimizer.step() #更新梯度

    if(i+1) % 20 ==0:
        print('Epoch[{}/{}],loss:{:.6f}'.format(i,5000,loss.data))


#4、模型评估
model.eval() #设置模型为评估模式，即预测模式
predict = model(x)
predict = predict.cpu().detach().numpy()
plt.scatter(x.cpu().detach().numpy(),y.cpu().detach().numpy(),c="r")
plt.plot(x.cpu().detach().numpy(),predict)
plt.show()