提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档
一、基本介绍
今天和大家分享的例子代码是用pytorch实现一个卷积神经网络实现图像识别，所用到的数据集是cifar10，是一个十分类的图像分类数据集，每个对象的所属类别为1类，总共类别为10类，输入图像数据的维度是[1，1，32，32].下面是具体的代码实现部分。
二、代码实现
1.了解数据
下载cifar10数据集，获得训练集、测试集；将其用dataloader进行封装；选一条数据打印出来看看。
import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
trainset=torchvision.datasets.CIFAR10(root='./data',train=True,download=True,transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True,num_workers=2)
testset=torchvision.datasets.CIFAR10(root='./data',train=False,download=True,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')
import matplotlib.pyplot as plt
#构建展示图片的函数
def imshow(img):
img=img/2+0.5
npimg=img.numpy()
plt.imshow(np.transfose(npimg,(1,2,0)))
plt.show()
#从数据迭代器中读取一张图片
dataiter=iter(trainloader)
images,labels=dataiter.next()
#展示图片
imshow(torchvision.utils.make_grid(images))
#打印标签
print(''.join('%5s'% classes[labels[j]] for j in range(4)))
2.构建网络、损失函数、优化器
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
#定义两个卷积层
self.conv1=nn.Conv2d(1,6,5)#有RGB的图片第一个参数channel应该为3,这里为了看size的变化才写1的
self.conv2=nn.Conv2d(6,16,5)
#定义池化层
self.pool=nn.MaxPool2d(2,2)
#定义三个全连接层
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=x.view(-1,16*5*5)
x=F.relu(self.fc1(x))
x=F.relu(self.fc2(x))
x=self.fc3(x)
return x
#input[1, 1, 32, 32]->conv1poll[1, 6, 14, 14]->conv2pool[1, 16, 5, 5]->view[1, 400]->linear1relu[1, 120]->linear2relu[1,84]->linear3relu[10]
net=Net()
print(net)
import torch.optim as optim
#定义损失函数，选用交叉熵损失函数
criterion=nn.CrossEntropyLoss()
#定义优化器，选用随机梯度下降的优化器
optimizer=optim.SGD(net.parameters(),lr=0.001,momentum=0.9)
3.训练模型
#编写训练集代码
for epoch in range(2):
running_loss=0.0
#按批次迭代训练模型
for i,data in enumerate(trainloader,0):
#从data中去除含有输入图像的张量inputs，标签张量labels
inputs,labels=data
#第一步梯度清零
optimizer.zero_grad()
#第二步将输入图像进入网络中，得到输出张量
outputs=net(inputs)
#计算损失值
loss=criterion(outputs,labels)
#进行反向传播和梯度更新
loss.backward()
optimizer.step()
#打印训练的信息
running_loss+=loss.item()
if (i+1)%2000==0:
print('[%d,%5d] loss:%.3f'%(epoch+1,i+1,running_loss/2000))
running_loss=0.0
print('Finished training')
4.模型保存与加载
#设定模型的保存位置
PATH='./cifar_net.pth'
#保存模型的状态字典
torch.save(net.state_dict(),PATH)#***********************************************
#在测试集中取出一个批次的数据，做图像和标签的展示
dataiter=iter(testloader)
images,labels=dataiter.next()
#打印原始图片
imshow(torchvision.utils.make_grid(images))
#打印真实的标签
print('GroundTruth',''.join('%5s'% classes[labels[j]] for j in range(4)))
#加载模型参数，在测试阶段
net.load_state_dict(torch.load(PATH))#*****************************************
#利用模型对图片进行预测
outputs=net(images)
#模型有10个类别的输出，选取其中概率最大的那个类别作为预测值
_,predicted=torch.max(outputs,1)
#打印预测标签
print('Predicted:',''.join('%5d'% classes[predicted[j]] for j in range(4)))
5.测试
#在这整个测试集上测试模型的准确率
correct=0
total=0
with torch.no_grad():
for data in testloader:
images,labels=data
outputs=net(images)
_,predicted=torch.max(outputs.data,1)
total+=labels.size(0)
correct+=(predicted==labels).sum().item()
print('Accuracy of the network on the 10000 test images:%d %%'%(100*correct/total))
#分别测试不同类别的模型准确率
class_correct=list(0. for i in range(10))
class_total=list(0. for i in range(10))
with torch.no_grad():
for data in testloader:
images,labels=data
outputs=net(images)
_,predicted=torch.max(outputs,1)
c=(predicted==labels).squeeze()
for i in range(4):
label=labels[i]
class_correct[label]+=c[i].item()
class_total[label]+=1
#打印不同类别的准确率
for i in range(10):
print('Accurary of %5s : %2d %%'% (classes[i],100*class_correct[i]/class_total[i]))
6.GPU
#在GPU上训练模型
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
#将模型转移到GPU上
net.to(device)
#将输入的图片张量和标签张量转移到GPU上
inputs,labels=data[0].to(device),data[1].to(device)
总结
今天最重要的部分就是模型的训练和测试啦，除此以外害很重要的是模型的保存和加载，这两个部分常常不引起重视，但自己做例子的时候会发现有很多问题，容易犯错，遂加*标以便及时巩固，加深印象。