pytorch如何使用GPU

在本文中,我将介绍简单如何使用GPU
pytorch是一个非常优秀的深度学习的框架,具有速度快,代码简洁,可读性强的优点。
我们使用pytorch做一个简单的回归。
首先准备数据

import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.autograd import  Variable
import torch.nn as nn
x = np.random.randn(1000, 1)*4
w = np.array([0.5,])
bias = -1.68

y_true = np.dot(x, w) + bias  #真实数据
y = y_true + np.random.randn(x.shape[0])#加噪声的数据
#我们需要使用x和y,以及y_true回归出w和bias

定义回归网络的类

class LinearRression(nn.Module):
    def __init__(self, input_size, out_size):
        super(LinearRression, self).__init__()
        self.x2o = nn.Linear(input_size, out_size)
    #初始化
    def forward(self, x):
        return self.x2o(x)
    #前向传递

接下来介绍将定义模型和优化器

batch_size = 10
model = LinearRression(1, 1)#回归模型
criterion = nn.MSELoss()  #损失函数
#调用cuda
model.cuda()
criterion.cuda()

optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
losses = []

下面就是训(练)练(丹)了

for i in range(epoches):
    loss = 0
    optimizer.zero_grad()#清空上一步的梯度
    idx = np.random.randint(x.shape[0], size=batch_size)
    batch_cpu = Variable(torch.from_numpy(x[idx])).float()
    batch = batch_cpu.cuda()#很重要

    target_cpu = Variable(torch.from_numpy(y[idx])).float()
    target = target_cpu.cuda()#很重要
    output = model.forward(batch)
    loss += criterion(output, target)
    loss.backward()
    optimizer.step()

    if (i +1)%10 == 0:
        print('Loss at epoch[%s]: %.3f' % (i, loss.data[0]))
    losses.append(loss.data[0])

plt.plot(losses, '-or')
plt.xlabel("Epoch")
plt.xlabel("Loss")

plt.show()

下面是训练结果

Loss at epoch[9]: 5.407
Loss at epoch[19]: 3.795
Loss at epoch[29]: 2.352
Loss at epoch[39]: 1.725
Loss at epoch[49]: 1.722
Loss at epoch[59]: 1.044
Loss at epoch[69]: 1.044
Loss at epoch[79]: 0.771
Loss at epoch[89]: 1.248
Loss at epoch[99]: 1.862

总结一下。要调用cuda执行代码需要一下步骤

model.cuda()
criterion.cuda()

以及

batch_cpu = Variable(torch.from_numpy(x[idx])).float()
    batch = batch_cpu.cuda()
    target_cpu = Variable(torch.from_numpy(y[idx])).float()
    target = target_cpu.cuda()

就是将模型和输入数据变为cuda执行的
,简直超级方便,良心推荐一波pytorch
最后,女神走一波

Logo

免费领 200 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐