第一次使用pytorch程序的gpu加速【实测成功】

仅作为操作记录，大佬请跳过。

背景

博主使用pytorch的cpu版本，对CIFAR10图像分类传送门进行了实现；
后将pytorch从cpu版本改装成gpu版本传送门，并安装cuda传送门，下面就是如何用gpu跑程序的问题了

解决

其实只需加四行代码（目的是定义cuda名字和输入输出值放到cuda），在原有的pytorch的cpu版本的代码上补写代码。

要添加的程序只有：

device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # add
print(device)  # add

net.to(device)  # add

inputs,labels=inputs.to(device),labels.to(device)  # add

————————————————————————————

可以在程序运行中，通过这种方法测试，判断是不是gpu在运行。

完整代码

原有的cpu版本：

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from time import time  # add add

transform=transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
trainset=torchvision.datasets.CIFAR10(root=r'E:\document\pytorchcifarclass',train=True,download=False,transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True,num_workers=2)

testset = torchvision.datasets.CIFAR10(root=r'E:\document\pytorchcifarclass', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
print('ok')

# network
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1=nn.Conv2d(3,6,5)
        self.pool=nn.MaxPool2d(2,2)
        self.conv2=nn.Conv2d(6,16,5)
        self.fc1=nn.Linear(16*5*5,120)
        self.fc2=nn.Linear(120,84)
        self.fc3=nn.Linear(84,10)

    def forward(self,x):
        x=self.pool(F.relu(self.conv1(x)))
        x=self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

if __name__=='__main__':
    net=Net()

    # loss and optim
    criterion=nn.CrossEntropyLoss()
    optimizer=optim.SGD(net.parameters(),lr=0.001,momentum=0.9)

    t1=time()  # add add

    # train
    for epoch in range(2):

        running_loss=0.0
        for i,data in enumerate(trainloader,0):

            inputs,labels=data

            optimizer.zero_grad()

            outputs=net(inputs)
            loss=criterion(outputs,labels)
            loss.backward()
            optimizer.step()

            running_loss+=loss.item()
            if i % 2000==1999:
                print('[%d,%5d] loss:%.3f'%(epoch+1,i+1,running_loss/2000))
                running_loss=0.0

    t2=time()  # add add
    t=t2-t1  # add add
    print('cpu耗时')
    print(t)  # add add

    print('Finished Training')

    # 保存网络
    # torch.save(net.state_dict(),'net_params.pkl')

经改写后的gpu版本：

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from time import time  # add add

transform=transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
trainset=torchvision.datasets.CIFAR10(root=r'E:\document\pytorchcifarclass',train=True,download=False,transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True,num_workers=2)

testset = torchvision.datasets.CIFAR10(root=r'E:\document\pytorchcifarclass', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
print('ok')

# network
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1=nn.Conv2d(3,6,5)
        self.pool=nn.MaxPool2d(2,2)
        self.conv2=nn.Conv2d(6,16,5)
        self.fc1=nn.Linear(16*5*5,120)
        self.fc2=nn.Linear(120,84)
        self.fc3=nn.Linear(84,10)

    def forward(self,x):
        x=self.pool(F.relu(self.conv1(x)))
        x=self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

if __name__=='__main__':
    device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # add
    print(device)  # add

    net=Net()

    net.to(device)  # add

    # loss and optim
    criterion=nn.CrossEntropyLoss()
    optimizer=optim.SGD(net.parameters(),lr=0.001,momentum=0.9)

    t1=time()  # add add

    # train
    for epoch in range(2):

        running_loss=0.0
        for i,data in enumerate(trainloader,0):

            inputs,labels=data

            inputs,labels=inputs.to(device),labels.to(device)  # add

            optimizer.zero_grad()

            outputs=net(inputs)
            loss=criterion(outputs,labels)
            loss.backward()
            optimizer.step()

            running_loss+=loss.item()
            if i % 2000==1999:
                print('[%d,%5d] loss:%.3f'%(epoch+1,i+1,running_loss/2000))
                running_loss=0.0

    t2=time()  # add add
    t=t2-t1  # add add
    print(t)  # add add

    print('Finished Training')

    # 保存网络
    # torch.save(net.state_dict(),'net_params.pkl')

注：未下载过CIFAR-10数据集的客友萌，需要将这两行的False改为True