初学者必看-pytorch框架下基于cifar-10数据集的图像分类模型

本博客基于pytorch对cifar10数据集实现分类模型。一些内容参考于互联网与GPT，如有问题，请大家及时评论指正。

目录

1、数据集导入

2、CNN网络搭建

3、训练

4、测试

5、结果

6、总代码

---

1、数据集导入

（1）通过torch的tranformer预处理数据：其中.Compose()函数表示将多个图片操作联合起来，而.ToTensor()则代表将图片格式转换成为tensor张量格式，.Normalize()代表归一化处理，其中最主要的做法就是：得到数据集的均值与标准差，再让每个样本减去均值，除以标准差，让图片从0~32转换到[0,1]区间内。

transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])

（2）数据集的加载与预处理：

用trainset来接受训练集，testset来接受测试集。用torch.utils.data中的DataLoader来load前面的trainset与testset，并分别用trainloader与testloader接收它们。其中batch_size=4,shuffle=True表示加载数据顺序被打乱。

trainset=torchvision.datasets.CIFAR10(root="E:/program_projects/NLP/torch.nn/cifar_model/data",train=True,download=True,transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True)

testset=torchvision.datasets.CIFAR10(root="E:/program_projects/NLP/torch.nn/cifar_model/data",train=False,download=True,transform=transform)
testloader=torch.utils.data.DataLoader(testset,batch_size=4,shuffle=False)

（3）用.to（）来调用GPU

用device来接收GPU作为训练设备。注意，GPU在代码中使用“cuda”而不是“gpu”

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

2、CNN网络搭建

因为我们使用的是pytorch，因此搭建cifarnet可以调用其ConV2d很方便构造出网络。本次定义中，使用了五层网络架构：

卷积层1（input_channel=3,output_channel=6,卷积核=5*5)-->

池化(padding=2,stride=2)-->

卷积层2（input_channel=6,output_channel=16,卷积核=5*5)-->

池化(padding=2,stride=2)-->

全连接层1-->全连接层2-->全连接层3

结构图如下图所示：

class cifarNet(nn.Module):
    def __init__(self):
        super(cifarNet,self).__init__()
        self.conv1=nn.Conv2d(3,6,5)
        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.view(-1,16*5*5)
        x=F.relu(self.fc1(x))
        x=F.relu(self.fc2(x))
        x=self.fc3(x)
        return x
cifarnet=cifarNet()
print(cifarnet)

3、训练

(1)选择优化器

使用torch.optim中的随机梯度下降算法SGD，momentum表示加入动量的SGD算法，可想象成一个有重量的小球滚向山谷，此时它的重量有利于其更快的到达山谷最低点，即损失函数极小点，此时它的自变量wi为我们寻找的最优权重。

optimozer=torch.optim.SGD(cifarnet.parameters(),lr=0.01,momentum=0.9)

(2)损失函数选择

分类模型最常选择交叉熵损失函数，这个在nn中自带。

criterion=nn.CrossEntropyLoss()

(3)样本在训练时，首先进行前向输入forward，再进行反向传播backward，但在这个环节，我们要注意在反向传播之前，必须要将之前的梯度置空，否则会在每次训练时积累梯度，导致更新的梯度是错误的。使用.zero_grad()置空梯度。其次，代码中使用loss.backward()时，本质上是反向求其梯度，而并没有真正的进行权重更新，我们还需要加入optimizer.step()来进行真正的更新！

cifarnet.to(device="cuda")
for epoch in range(5):
    running_loss=0.0
    for i,data in enumerate(trainloader,0):# 对于载入的训练数据集中的每个样本 data
        inputs,labels=data
        inputs, labels = data[0].to(device), data[1].to(device)# 将数据传入GPU设备 并用inputs labels去接受
        optimizer.zero_grad()# 梯度清空
        outputs=cifarnet(inputs) # 将数据传入cifarnet的CNN网络 用outputs接受
        loss=criterion(outputs,labels)# 得到 预测结果与真实labels之间的损失
        loss.backward()# 反向计算梯度
        optimizer.step()# 梯度更新

        running_loss+=loss.item()# 累加running loss
        if(i+1)%2000==0: # 每2000个样本记录一次loss
            print('[%d,%5d0 loss:%.3f'% (epoch+1,i+1,running_loss/2000))
            running_loss=0.0# 将running loss置零 重新累加

print('finished training')
Path='./cifar_net.pth'
torch.save(cifarnet.state_dict(),Path)# 保存训练模型的状态字典
print("cifarnet saved in ./cifar_net.pth")

4、测试

（1）用总体测试集测试训练好的cifarnet模型

用.load_state_dict(torch.load(Path))来解码刚保存的模型状态字典。

correct=0.0# 正确数量
total=0.0# 总数
with torch.no_grad():
     for data in testloader:#对于测试集的data
         images,labels=data[0].to(device), data[1].to(device)# 传入GPU
         outputs=cifarnet(images)# image 输入网络
         _,prediction=torch.max(outputs.data,1)# 贪心算法得到概率最大的一个类
         total +=labels.size(0)# 全部测试样本的数量
         correct +=(prediction==labels).sum().item()# 对比prediction与labels是否相同，相同为1，不同为0，为1的累加在一起，即得到正确预测的数量
print('Accuracy of the cifarnet on the 10000 test images:%d,%% ' % (100*correct/total))

（2）不同类别的准确率

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[0].to(device), data[1].to(device)
        outputs=cifarnet(images)
        _,prediction=torch.max(outputs,1)
        c=(prediction==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('Accuracy of %5s: %2d %%' % (classes[i],100*class_correct[i]/class_total[i]))

5、结果

我设定训练了5次之后的结果如下：

6、总代码

import numpy as np
import torch
import  torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils import data
from torch.utils.data import DataLoader

import numpy
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim
from cifar_net import cifarnet


transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
trainset=torchvision.datasets.CIFAR10(root="E:/program_projects/NLP/torch.nn/cifar_model/data",train=True,download=True,transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True)

testset=torchvision.datasets.CIFAR10(root="E:/program_projects/NLP/torch.nn/cifar_model/data",train=False,download=True,transform=transform)
testloader=torch.utils.data.DataLoader(testset,batch_size=4,shuffle=False)

classes=("plane","car","bird","cat","deer","dog","frog","horse","ship","truck")
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# cifarnet CNN模型
class cifarNet(nn.Module):
    def __init__(self):
        super(cifarNet,self).__init__()
        self.conv1=nn.Conv2d(3,6,5)
        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.view(-1,16*5*5)
        x=F.relu(self.fc1(x))
        x=F.relu(self.fc2(x))
        x=self.fc3(x)
        return x
cifarnet=cifarNet()
print(cifarnet)

# 优化器选择、损失函数定义
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(cifarnet.parameters(),lr=0.01,momentum=0.9)

# 训练
cifarnet.to(device="cuda")
for epoch in range(5):
    running_loss=0.0
    for i,data in enumerate(trainloader,0):# 对于载入的训练数据集中的每个样本 data
        inputs,labels=data
        inputs, labels = data[0].to(device), data[1].to(device)# 将数据传入GPU设备 并用inputs labels去接收
        optimizer.zero_grad()# 梯度清空
        outputs=cifarnet(inputs) # 将数据传入cifarnet的CNN网络 用outputs接受
        loss=criterion(outputs,labels)# 得到 预测结果与真实labels之间的损失
        loss.backward()# 反向计算梯度
        optimizer.step()# 梯度更新

        running_loss+=loss.item()# 累加running loss
        if(i+1)%2000==0: # 每2000个样本记录一次loss
            print('[%d,%5d0 loss:%.3f'% (epoch+1,i+1,running_loss/2000))
            running_loss=0.0# 将running loss置零 重新累加

print('finished training')
Path='./cifar_net.pth'
torch.save(cifarnet.state_dict(),Path)# 保存训练模型的状态字典
print("cifarnet saved in ./cifar_net.pth")

cifarnet.load_state_dict(torch.load(Path))# 解码已保存的cifarnet状态字典

# correct=0.0# 正确数量
# total=0.0# 总数
# with torch.no_grad():
#     for data in testloader:#对于测试集的data
#         images,labels=data[0].to(device), data[1].to(device)
#         outputs=cifarnet(images)
#         _,prediction=torch.max(outputs.data,1)
#         total +=labels.size(0)
#         correct +=(prediction==labels).sum().item()
# print('Accuracy of the cifarnet 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[0].to(device), data[1].to(device)
        outputs=cifarnet(images)
        #outputs=output.to(device)
        _,prediction=torch.max(outputs,1)
        c=(prediction==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('Accuracy of %5s: %2d %%' % (classes[i],100*class_correct[i]/class_total[i]))