Pytorch图像分类从模型自定义到测试

前面已跟大家介绍了 Caffe 和 TensorFlow，现在说说 Pytorch，集齐三大主流框架以后方便召唤模型。

1 什么是 Pytorch

一句话总结 Pytorch = Python + Torch。

Torch 是纽约大学的一个机器学习开源框架，几年前在学术界非常流行，包括 Lecun 等大佬都在使用。但是由于使用的是一种绝大部分人绝对没有听过的 Lua 语言，导致很多人都被吓退。后来随着 Python 的生态越来越完善，Facebook 人工智能研究院推出了 Pytorch 并开源。Pytorch 不是简单的封装 Torch 并提供 Python 接口，而是对 Tensor 以上的所有代码进行了重构，同 TensorFlow 一样，增加了自动求导。

后来 Caffe2 全部并入 Pytorch，如今已经成为了非常流行的框架。很多最新的研究如风格化、GAN 等大多数采用 Pytorch 源码，这也是我们必须要讲解它的原因。

Pytorch 有什么特点呢

（1）动态图计算。前面说过，TensorFlow 是采用静态图的，先定义好图，然后在 session 中运算。图一旦定义好后，是不能随意修改的。当然了，现在 TensorFlow 也引入了动态图机制 Eager Execution，只是不如 Pytorch 直观。那动态图有什么好处呢？TensorFlow 要查看变量结果，必须在 sess 中，sess 的角色是 C 语言的执行，而之前的图定义是编译。而 Pytorch 就好像是脚本语言，随时随地修改，随处 debug，没有一个类似编译的过程，这比 TensorFlow 要灵活很多。

（2）简单。TensorFlow 的学习成本真的不低，对于新手来说，Tensor、Variable、Session 等概念充斥，数据读取接口频繁更新，tf.nn、tf.layers、tf.contrib 各自重复，Pytorch 则是从 Tensor 到 Variable 再到 nn.Module，分别就是从数据张量到网络的抽象层次的递进。有人调侃 TensorFlow 的设计是“make it complicated”，那么 Pytorch 的设计就是“keep it simple”。

Pytorch 重要概念

（1）Tensor

这几大框架都有基本的数据结构，Caffe 是 blob，TensorFlow 和 Pytorch 都是 Tensor，都是高维数组。Pytorch 中的 Tensor 使用与 Numpy 的数组非常相似，两者可以互转且共享内存。

（2）Variable

Variable 封装 Tensor，然后提供反向传播，自动计算梯度。Variable 实际上包含3个属性，data，即 Tensor 内容；Grad，是与 data 对应的梯度；grad_fn，计算反向传播的函数。当然了，就在不到一个月前Pytorch 0.4.0已经合并了Variable和Tensor。

（3）Nn.module

抽象好的网络数据结构，可以表示为网络的一层，也可以表示为一个网络结构。在实际使用过程中，经常会定义自己的网络，并继承 nn.Module。具体的使用，我们看下面的网络定义吧。

2 Pytorch 训练

安装咱们就不说了，接下来的任务就是开始训练模型。训练模型包括数据准备、模型定义、结果保存与分析。

2.1 数据准备与读取

前面已经介绍了 Caffe 和 TensorFlow 的数据读取，两者的输入都是图片 list，但是读取操作过程差异非常大，Pytorch 与这两个又有很大的差异。这一次，直接利用文件夹作为输入，这是 Pytorch 更加方便的做法。

数据读取的完整代码如下：

    data_dir = '../../../../datas/head/'    
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomSizedCrop(48),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
        ]),
        'val': transforms.Compose([
            transforms.Scale(64),
            transforms.CenterCrop(48),
            transforms.ToTensor(),
            transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
        ]),
    }
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                              data_transforms[x]) for x in ['train', 'val']}
    dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
                                                 batch_size=16,
                                                 shuffle=True,
                                                 num_workers=4) for x in ['train', 'val']}

下面一个一个解释，完整代码请移步 Git 工程。

（1）datasets.ImageFolder

Pytorch 的 torchvision 模块中提供了一个 dataset 包，它包含了一些基本的数据集如 mnist、coco、imagenet 和一个通用的数据加载器 ImageFolder，可以制定train和val目录，而其中的每一个目录则分为不同类别的文件夹。

具体的请到 Git 工程中查看。

├── train
│   ├── 0
│   ├── 1
└── val
    ├── 0
    ├── 1

imagefolder 有3个成员变量。

• self.classes：用一个 list 保存类名，就是文件夹的名字。
• self.class_to_idx：类名对应的索引，可以理解为 0、1、2、3 等。
• self.imgs：保存（imgpath，class），是图片和类别的数组。

不同文件夹下的图，会被当作不同的类，天生就用于图像分类任务。

（2）Transforms

这一点跟 Caffe 非常类似，就是定义了一系列数据集的预处理和增强操作。到此，数据接口就定义完毕了，接下来在训练代码中看如何使用迭代器进行数据读取就可以了，包括 scale、减均值等。

（3）torch.utils.data.DataLoader

这就是创建了一个 batch，生成真正网络的输入。关于更多 Pytorch 的数据读取方法，请移步知乎专栏和公众号，链接在前面的课程中有给出。

2.2 模型定义

如下：

    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import numpy as np
    class simpleconv3(nn.Module):`
    def __init__(self):
        super(simpleconv3,self).__init__()
        self.conv1 = nn.Conv2d(3, 12, 3, 2)
        self.bn1 = nn.BatchNorm2d(12)
        self.conv2 = nn.Conv2d(12, 24, 3, 2)
        self.bn2 = nn.BatchNorm2d(24)
        self.conv3 = nn.Conv2d(24, 48, 3, 2)
        self.bn3 = nn.BatchNorm2d(48)
        self.fc1 = nn.Linear(48 * 5 * 5 , 1200)
        self.fc2 = nn.Linear(1200 , 128)
        self.fc3 = nn.Linear(128 , 2)
    def forward(self , x):
        x = F.relu(self.bn1(self.conv1(x)))
        #print "bn1 shape",x.shape        
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = x.view(-1 , 48 * 5 * 5) 
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

这三节课的任务，都是采用一个简单的 3 层卷积 + 2 层全连接层的网络结构。根据上面的网络结构的定义，需要做以下事情。

（1）simpleconv3(nn.Module)

继承 nn.Module，前面已经说过，Pytorch 的网络层是包含在 nn.Module 里，所以所有的网络定义，都需要继承该网络层。

并实现 super 方法，如下：

super(simpleconv3,self).__init__()

这个，就当作一个标准，执行就可以了。

（2）网络结构的定义

都在 nn 包里，举例说明：

torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)

完整的接口如上，定义的第一个卷积层如下：

nn.Conv2d(3, 12, 3, 2)

即输入通道为3，输出通道为12，卷积核大小为3，stride=2，其他的层就不一一介绍了，大家可以自己去看 nn 的 API。

（3）forward

backward 方法不需要自己实现，但是 forward 函数是必须要自己实现的，从上面可以看出，forward 函数也是非常简单，串接各个网络层就可以了。

对比 Caffe 和 TensorFlow 可以看出，Pytorch 的网络定义更加简单，初始化方法都没有显示出现，因为 Pytorch 已经提供了默认初始化。

如果我们想实现自己的初始化，可以这么做：

init.xavier_uniform(self.conv1.weight)init.constant(self.conv1.bias, 0.1)

它会对 conv1 的权重和偏置进行初始化。如果要对所有 conv 层使用 xavier 初始化呢？可以定义一个函数：

def weights_init(m):    
    if isinstance(m, nn.Conv2d):
        xavier(m.weight.data)
        xavier(m.bias.data)  
    net = Net()   
    net.apply(weights_init)

3 模型训练

网络定义和数据加载都定义好之后，就可以进行训练了，老规矩先上代码：

    def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
        for epoch in range(num_epochs):
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            for phase in ['train', 'val']:
                if phase == 'train':
                    scheduler.step()
                    model.train(True)  
                else:
                    model.train(False)  
                    running_loss = 0.0                running_corrects = 0.0
                for data in dataloders[phase]:
                    inputs, labels = data
                    if use_gpu:
                        inputs = Variable(inputs.cuda())
                        labels = Variable(labels.cuda())
                    else:
                        inputs, labels = Variable(inputs), Variable(labels)

                    optimizer.zero_grad()
                    outputs = model(inputs)
                    _, preds = torch.max(outputs.data, 1)
                    loss = criterion(outputs, labels)
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                    running_loss += loss.data.item()
                    running_corrects += torch.sum(preds == labels).item()

                epoch_loss = running_loss / dataset_sizes[phase]
                epoch_acc = running_corrects / dataset_sizes[phase]

                if phase == 'train':
                    writer.add_scalar('data/trainloss', epoch_loss, epoch)
                    writer.add_scalar('data/trainacc', epoch_acc, epoch)
                else:
                    writer.add_scalar('data/valloss', epoch_loss, epoch)
                    writer.add_scalar('data/valacc', epoch_acc, epoch)

                print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))
        writer.export_scalars_to_json("./all_scalars.json")
        writer.close()
        return model

分析一下上面的代码，外层循环是 epoches，然后利用 for data in dataloders[phase] 循环取一个 epoch 的数据，并塞入 variable，送入 model。需要注意的是，每一次 forward 要将梯度清零，即 optimizer.zero_grad()，因为梯度会记录前一次的状态，然后计算 loss，反向传播。

loss.backward()
optimizer.step()

下面可以分别得到预测结果和 loss，每一次 epoch 完成计算。

epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects / dataset_sizes[phase]
      
_, preds = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)

4 可视化

可视化是非常重要的，鉴于 TensorFlow 的可视化非常方便，我们选择了一个开源工具包，tensorboardx，安装方法为 pip install tensorboardx，使用非常简单。

第一步，引入包定义创建：

from tensorboardX import SummaryWriter
writer = SummaryWriter()

第二步，记录变量，如 train 阶段的 loss，writer.add_scalar('data/trainloss', epoch_loss, epoch)。

按照以上操作就完成了，完整代码可以看上面的 Git 项目，我们看看训练中的记录。Loss 和 acc 的曲线图如下：

 

网络的收敛没有 Caffe 和 TensorFlow 好，大家可以自己去调试调试参数了，随便折腾吧。

5 Pytorch 测试

上面已经训练好了模型，接下来的目标就是要用它来做 inference 了，同样，给出代码。

    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.optim import lr_scheduler
    from torch.autograd import Variable
    import torchvision
    from torchvision import datasets, models, transforms
    import time
    import os
    from PIL import Image
    import sys
    import torch.nn.functional as F

    from net import simpleconv3
    data_transforms =  transforms.Compose([
                transforms.Resize(48),
                transforms.ToTensor(),
                transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])])

    net = simpleconv3()
    net.eval()
    modelpath = sys.argv[1]
    net.load_state_dict(torch.load(modelpath,map_location=lambda storage,loc: storage))

    imagepath = sys.argv[2]
    image = Image.open(imagepath)
    imgblob = data_transforms(image).unsqueeze(0)
    imgblob = Variable(imgblob)

    torch.no_grad()

    predict = F.softmax(net(imgblob))
    print(predict)

从上面的代码可知，做了几件事：

• 定义网络并使用 torch.load 和load_state_dict载入模型。
• 用 PIL 的 Image 包读取图片，这里没有用 OpenCV，因为 Pytorch 默认的图片读取工具就是 PIL 的 Image，它会将图片按照 RGB 的格式，归一化到 0～1 之间。读取图片之后，必须转化为 Variable 变量。
• evaluation 的时候，必须设置torch.no_grad()，然后就可以调用 softmax 函数得到结果了。

总结：本节讲了如何用 Pytorch 完成一个分类任务，并学习了可视化以及使用训练好的模型做测试。