第一次使用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
