【python】pytorch包（第三章）GPU上运行代码

一、实现流程

step 1. 判断本机是否支持GPU

torch.cuda.is_available()
结果为True即可行，否则不可行

step 2. 设定运行环境、数据类型为GPU

#设定环境为GPU
torch.device( "cuda:0" )
#设定模型，使适用于GPU
model.to(device)
#将tensor变量更改为GPU适用的类型
tensor.to(device)

step 3. 运行代码，计算结果

假定得到结果的变量命名为 pred_y

step 4. 将结果的数据类型转为可读类型

# 转为张量数据类型
pred_y_tensor = pred_y.detach()
# 转为numpy数据类型
pred_y_np = pred_y.detach().numpy()

二、 代码案例

import torch
# 判断本机是否支持GPU
torch.cuda.is_available()
#设定环境为GPU
device  = torch.device( "cuda:0" )

## Part 0.  数据准备
import torch
# 准备数据
x = torch.rand([100,1]) #500个数据，每个数据都是一维数据
y = x*3 + 0.8
#数据类型转换为GPU适用类型
x = x.to(device)
y = y.to(device)


## Part 1. 定义模型逻辑
from torch import nn
class Lr(nn.Module):
	def __init__(self):
		super(Lr,self).__init__() #继承父类init的参数
		self.linear = nn.Linear(1,1)
	def forward(self,x): #模型输入x，输出out
		out = self.linear(x)
		return out
#本段用于定义模型的逻辑关系


## Part 2. 模型实例化（设定模型参数）
from torch import optim
#调用模型基底、设定Loss函数、设定优化器
model = Lr()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(),lr=1e-3)
#模型转化为GPU类型
model.to(device)


##Part 3. 训练模型
#训练模型
training_times = int(input("输入模型训练次数："))
for i in range(training_times):
	out = model(x) #模型当前状态的预测值
	loss = criterion(y,out) #损失值
	optimizer.zero_grad() #梯度归零
	loss.backward() #计算梯度
	optimizer.step() #更新梯度
	if (i-1)%20 == 0:
		print('Epoch[%d/%d],loss:%.6f' % (i,training_times,loss.data))


##Part 4. 预测与评估
model.eval() #使模型进入评估模式/预测模式
#该模式下，输入的新数据不会让模型受到新的训练
pred_y = model(x) #输入自变量x，获得预测值pred_y
pred_y = pred_y.detach.numpy() #转darray类型

#可将 (x,y)点集以及(pred_y,x)预测曲线进行绘图，可视化数据拟合情况
import numpy
from matplotlib import pyplot as plt
plt.scatter(x.data.numpy(),y.data.numpy(),c='red')
plt.plot(x.data.numpy(),pred_y,c='green')
plt.show()
print(loss.data.item())