pytorch学习【训练分类器】

说明：此博客内容来源于pytorch官方文档英文版，链接：pytorch.org

对于计算机视觉，我们已经创建了一个名为torchvision的包，该包含有支持加载类似Imagenet、CIFAR10，MNIST等公共数据集的数据加载模块torchvision.datasets和支持加载图像数据转换模块torch.utils.data.DataLoader.
对于本教程，我们使用公共数据集CIFAR10，它包含10个类别：airplane、automobile、bird、cat、deer、dog、frog、horse、ship、truck。CIFAR10的图像尺寸在33232，其中3为RGB三个颜色通道，每个通道内图像尺寸为32*32.

训练一个图像分类器的步骤：1.使用torchvision加载并且归一化CIFAR10的训练集和测试集；2.定义一个卷积神经网络；3.定义一个损失函数；4.在训练样本上训练神经网络；5.在测试样本上测试网络。

1.加载数据并归一化

import torch
import torchvision
import torchvision.transforms as transforms

#transforms.ToTensor()会将像素值转换为范围为[0,1]的tensor。Normalize两个参数分别为均值和方差，将像素值减掉均值后除以标准差差，得到标准化像素值范围为[-1,1]
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
#加载训练集。root为路径；train为True表示该数据集为训练集；download为True表示下载数据集，由于本人数据集已下载，因此设置为False。
trainset = torchvision.datasets.CIFAR10(root='.data',train=True,download=False,transform=transform)
#shuffle=True表示打乱数据集顺序，一般设置为True
trainloader = torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True)
testset = torchvision.datasets.CIFAR10(root='.data',train=False,download=False,transform=transform)
testloader=torch.utils.data.DataLoader(testset,batch_size=10,shuffle=True)

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

显示一下训练集

import numpy as np
import matplotlib.pyplot as plt

def imshow(img):
    #首先去归一化，因为前面对图像进行了归一化，即减掉均值后除以标准差（均值标准差均为0.5），因此去归一化需要乘以0.5再加0.5.
    img = img/2+0.5
    npimg = img.numpy()
    #由于输入图像格式为（C，W,H）形式，而numpy可以读取的形式为（W,H,C），因此需要将三个维度顺序转换一下
    plt.imshow(np.transpose(npimg,(1,2,0)))
    plt.show()
#定义一个迭代器，每次运行会显示下一组图像
detaiter = iter(trainloader)
images,labels=detaiter.next()

imshow(torchvision.utils.make_grid(images))
print(''.join('%5s'%classes[labels[j]] for j in range(4)))

2.定义神经网络

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        #定义两个卷积层，Conv2d为二维卷积。3，6，5分别为输入通道，输出通道，卷积核大小
        self.conv1 = nn.Conv2d(3,6,5)
        self.conv2 = nn.Conv2d(6,16,5)
        #定义三个全连接层，两个参数分别为输入单元数和输出单元数，由于CIFAR10有10中类别，因此最后一层输出单元个数为10。
        self.fc1 = nn.Linear(16*5*5,120)
        self.fc2 = nn.Linear(120,84)
        self.fc3 = nn.Linear(84,10)
    #定义好网络结构，开始前向传播
    def forward(self,x):
        #F.max_pool2d表示二维最大池化，（2，2）表示filter大小为2，如果filter为矩形，可以省略一个2。F.relu为relu激活函数
        x = F.max_pool2d(F.relu(self.conv1(x)),(2,2))
        x = F.max_pool2d(F.relu(self.conv2(x)),2)
        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
net = Net()

3.定义损失函数和优化器

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

4.训练神经网络

for epoch in range(2):
    running_loss = 0
    #enumerate两个参数(序列，起始位置）,返回值为枚举对象
    for i,data in enumerate(trainloader,0):
        inputs,labels = data
        #由于pytorch计算梯度会累积，因此每次反向传播之前需要清空梯度
        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
print("finish training")

我们可以保存已训练好的模型

PATH='./cifar_net.pth'
torch.save(net.state_dict(),PATH)
#读取已保存的模型
net = Net()
net.load_state_dict(torch.load(PATH))

5.测试数据集

net.eval()
with torch.no_grad():
    total_correct = 0
    total_num = 0
    for x,label in testloader:
        out = net(x)
        pred = out.argmax(dim=1)
        total_correct += torch.eq(pred,label).float().sum().item()
        total_num += x.size(0)

    acc = total_correct /total_num
    print(epoch,acc)

测试每一类的分类精度

class_correct = list(0. for i in range(10))  #记住这个表达，生成10个0的列表
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('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))