Pytorch 迁移学习:从 ImageNet 到 Caltech 101

• Report
ABC
本文属初学作品,其时仍在学习、探索;高人不必读。

从头训练模型,耗时耗力,往往做了大量的重复工作;而在预训练的模型基础上实施迁移学习,能够很快的使已经学习到丰富参数的网络适配具体问题,节省时间,并允许我们在模型上尝试更多方面的工作,而不是局限于训练本身。本项目参考一篇教程制作。

本文所用代码已发布在码云上:xjtu-blacksmith/transfer-learning-beginner

项目总述

采用由 Pytorch 提供的预训练模型,在新的数据集 Caltech 101 上进行迁移学习。在此过程中了解:

  • 如何冻结部分参数,只训练网络一部;
  • 如何在已有网络基础上附加新结构;
  • 尝试 grad_cam,可视化模型预测结果。

同时,在此次训练过程中,更多地探索以下方面的内容:

  • 分类结果检查;
  • 工具实用化(允许用脚本加载、预测任意图片);
  • 训练、分类结果的统计。

数据集准备

数据集划分

由于 Caltech 101 数据集中并无训练集、测试集的划分,该工作需要自行完成。根据教程的建议,设置划分比例为:训练集占 50%,测试集与验证集各占 25%。为此,采用 sklearn.model_selection 模块中的 train_test_split 方法,对数据集做两次划分,以得到三个数据集。

以下展示划分所用的代码,其通过处理原始数据集中的分类与图片名称、调用 os 模块的文件处理方法移动文件来完成。该代码在数据集顶级目录下执行。

import os
from os import path

import numpy as np
import progressbar
from sklearn.model_selection import train_test_split

data_dir = path.join('data', 'caltech-101')  # where the raw data are
categories = os.listdir(data_dir)  # get 101 categories

bar = progressbar.ProgressBar()  # set a default progressbar

for i in bar(range(len(categories))):

    category = categories[i]  # fetch a specific category
    cat_dir = path.join(data_dir, category)

    images = os.listdir(cat_dir)  # get all images under the category

    # images split by their names
    images, images_test = train_test_split(images, test_size=0.25)
    images_train, images_val = train_test_split(images, test_size=0.33)
    image_sets = images_train, images_test, images_val
    labels = 'train', 'test', 'val'

    # move to corresponding folders
    for image_set, label in zip(image_sets, labels):
        dst_folder = path.join(data_dir, label, category)  # create folder
        os.makedirs(dst_folder)
        for image in image_set:
            # rename ../cat/xxx.jpg to ../label/cat/xxx.jpg
            src_dir = path.join(cat_dir, image)
            dst_dir = path.join(dst_folder, image)
            os.rename(src_dir, dst_dir)  # move
    
    os.rmdir(cat_dir)  # remove empty folder

执行程序,本机上几秒钟即告终。打开数据集目录,初步观察,可见基本实现了既定要求。图 1 展示了三个文件夹在资源管理器中的情况。

图 1. 数据集划分后的目录

RGB 均值估算

根据在 ImageNet 上训练的经验,事先估计数据集中图片的 RGB 通道均值(及方差)是有利于训练的,估计参数可用于图形的标准化。之前未做过相关工作,参考了已有的方法(见这个帖子),通过如下函数在数据集上估计 RGB 均值与方差:

def compute_rgb_mean_std(data_set='train', batch_size=4):
    _, loader = build_data(data_set, batch_size, transform=False)
    mean, std, sample_num = 0., 0., 0.
    for data, _ in loader:
        batch_num = data.size(0)
        data = data.view(batch_num, data.size(1), -1)
        mean += data.mean(2).sum(0)
        std += data.std(2).sum(0)
        sample_num += batch_num
    
    mean /= sample_num
    std /= sample_num
    return mean, std

在 Caltech101 的 train 子集上运行该函数,获得的结果为 mean=[0.5429, 0.5263, 0.4994]std=[0.2422, 0.2392, 0.2406]。这些数据可以直接代入到后面的数据增强环节之中。

数据增强

在之前的训练之中,采用的数据增强方式比较单一,一般只是水平翻转和小范围的裁切。此次训练中,参考教程中的程序,对训练集采用了以下几项增强措施:

  • 随机裁切(80% – 100%),最终放缩为 256x256 的图片;
  • 随机水平翻转(概率为 50%);
  • 随机旋转,角度范围为 -15° – 15°;
  • 随机的颜色调整,采用默认参数、概率;
  • 中心裁切至 224x224 的尺寸;
  • 转为 Tensor,最终按上面估算的 RGB 均值与方差对这个 Tensor 进行标准化。

在以上增强措施中,最初的尺寸是 256x256,而最终的尺寸是 224x224;这样的做法,主要是为了使得图片在经过小角度的随即旋转后,能通过裁切去除旋转后多出的三角形空隙,并调整至各个标准模型的输入尺寸。

网络修改及训练参数

本次迁移学习,采用预训练的 VGG16 网络进行。Pytorch 中提供了下载预训练模型的模块,只要在创建网络时声明 pretrain=True 即可将预训练模型文件缓存至本地;问题在于,如何将预训练网络适配到本问题中。

可以创建一个 VGG16 的对象,通过 named_modules 属性检查其模块构成:

>>> from torchvision.models import vgg16
>>> net = vgg16()
>>> print(net.named_modules)
<bound method Module.named_modules of VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)>

可以看到,Pytorch 中定义的 VGG16 结构可分为三块:featuresavgpoolclassifier。显然,只有最终的分类层 classifier 是需要我们改动的,具体而言是最后的 out_features;考虑到从 4096 个 in_features 一跃输出至 Caltech101 数据集的 101 个 out_features 有些突兀,中间还可以再加上一个全连接层过渡,故直接将 classifier[6] 对应的线性层拓展为如下形式:

net.classifier[6] = nn.Sequential(
    nn.Linear(4096, 256),
    nn.ReLU(inplace=True),
    nn.Dropout(0.4),
    nn.Linear(256, CLASS_NUM),
)

其中 CLASS_NUM=101。这样,VGG16 就被改造为适于 Caltech101 数据集的模型了。为保证主干程序的简洁,将以上的步骤(创建对象、获取预训练参数、修改网络分类层)封装到一个函数中,由其获取到的网络对象命名为 vgg16_101

数据集、网络等都已准备停当,现在只需要配置优化器、代价函数了。本次训练中:

  • 优化器选择为 Adam,初始学习率设为 0.01,其他参数均取默认值;
  • 代价函数选择交叉熵,注意上面的 vgg16_101 网络中没有附加 LogSoftmax 层(否则代价函数应该改为 NLLLoss)。
  • 学习率衰减方式,沿袭上一次在 Tiny-ImageNet 上的做法:当验证集上准确率下降时,将学习率下调至原来的五分之一。根据上次获得的经验,在本次训练中一旦学习率下调总计三次即停止训练,不再要求这三次下调连续。这样可能不能达到最好的训练效果,但能在时间成本与训练效果方面达成更好的平衡。

训练结果与可视化

本次训练中,在之前的训练代码框架上进一步做了优化,将各个部分的代码整理至独立的函数、代码文件中;如训练部分整理至 train.py,评估部分整理至 val.py(其实此处命名有误,应该命名为 eval.py),inference、可视化等工作也有单独的代码、函数可供调用。在训练主干 main.py 文件中,现在只剩下这样几段代码:

from train import train_net
from model import vgg16_101

if __name__ == '__main__':
    net = vgg16_101(pretrained=True)
    train_net(net, lr=.01, epoch_num=30, msg='vgg16')

当然,这个 main.py 文件中只调用了训练的函数,其他函数未容纳其中。运行这段代码,即开始了训练,日志文件、数据文件根据训练情况不断被保存至 output 目录中。在笔记本上,训练不到一小时即可完成;在调试各处功能时反复执行过几次,这里仅以最终一次训练得到的模型做分析。

训练进度

与从头开始的训练不同,迁移学习中预训练过的模型已有比较良好的特征提取能力,故在其上误差下降、准确率上升都非常迅速。在 Caltech101 上,VGG16 网络经过 16 轮即达成了训练终止条件,训练结束时其在验证集上的准确率为 71.55%。

从训练输出的数据文件,可以绘出网络的误差下降与准确率变化情况,如图 2、图 3 所示。与从头训练的曲线相较,迁移训练过程中的误差曲线下降势头虽猛,却伴随着较大幅度的波动,同时也并未在学习率衰减的节点上有任何明显的转折;而准确率曲线中,训练集与测试集上的误差水平基本一致,表明此过程中没有出现明显的过拟合现象(这可能也与数据集本身特性有关)。

图 2. 训练误差曲线

图 3. 训练集与验证集上准确率曲线

准确率评估

在预先划分出来的测试集上,对训练得到的模型进评估,得到的准确率如表 1 所示。其中的两个模型,分别是验证集上准确率最高时所保存的模型,与训练结束时保存的模型(根据训练终止的条件,最终的模型自然不是准确率最高的)。

表 1. 模型在测试集上准确率

模型 Top-1 准确率 Top-5 准确率
最佳准确率模型 70.76% 84.81%
最终模型 70.67% 85.13%

两个模型的准确率大抵相近,不分伯仲。以上数据从侧面反映:在当前的学习率下降、训练终止策略下,模型的训练「潜力」已经基本发挥出来了。(但是,这并不是说模型已达到最优结果,如果改变策略可能会做得更好。)

陌生图像预测

为了检测模型对单张陌生(数据集外)图片的识别能力,从网络上准备了 6 张图片,如图 4 所示。其标签分别为:手风琴(accordion)、大炮(cannon)、吊扇(ceiling_fan)、小龙虾(crayfish)、椅子(chair)、阴阳(yin_yang)。

图 4. 陌生图片预览

将图片载入到程序之中后,将其处理成符合标准输入格式的、类于 minibatch 的形式(有四个维度的 Tensor),经模型预测得到结果。仿照上次在 TinyImageNet 上的训练,引入翻译模块将图片的英文标签翻译为中文;在此之外,还将网络对 Top-5 分类输出的数值经过 Softmax 处理,得到模型对各个分类预测的概率值,一并输出。在以上 6 张图片构成的目录中运行程序,得到的预测结果如表 2 所示。

表 2. 单张图片预测结果

序号 标签 预测结果 概率
1 手风琴 手风琴, 看, 手机, 三角钢琴, 面 99.96%, 0.04%, 0.00%, 0.00%, 0.00%
2 大炮 灯台, 鹦鹉螺, 阴阳, 订书机, 吊灯 21.88%, 16.93%, 10.75%, 5.73%, 4.54%
3 吊扇 灯台, 留声机, 灯, 面, 佛 36.13%, 35.69%, 16.74%, 8.46%, 1.42%
4 小龙虾 小龙虾, 蝎, 上, 螃蟹, 蜉蝣 69.69%, 15.56%, 6.73%, 3.17%, 1.49%
5 椅子 温莎椅, 椅子, 灯, 灯台, 笔记本电脑 53.28%, 46.71%, 0.01%, 0.01%, 0.00%
6 阴阳 阴阳, 熊猫, 海狸, 猬, 鹦鹉螺 100.00%, 0.00%, 0.00%, 0.00%, 0.00%

可以看到,模型在 1、4、6 图片中获得了正确的预测结果,在 5 号图片中正确预测结果位列第二,而在 2、3 号图片上预测结果错误。从概率数据可以进一步看出,模型在正确的预测中总是相当肯定(概率值很高),而在错误的预测中对猜测的各个分类都无完全的信心。这些结果都说明,模型的训练基本正常;在一些图片中,其已经积累了丰富的特征提取能力,而对另外一些图片则尚感陌生。考虑到数据集的容量并不很大(训练集不过四千余张图片,平均每个分类四十张图而已),取得这样的结果已经不错了。

利用 Grad-CAM 进行可视化

为了进一步弄清模型作出预测的依据,以上面 6 张陌生图片为预测对象,在模型中采用 Grad-CAM 方法进行可视化。Grad-CAM 的原理在提出其的论文中已经详细说明,粗略来说是:将图片馈入模型后,反向传播获得模型在指定层(往往取最后一个卷积层)上的梯度,对其求和、池化后将其作为对应分类的权重,并对应地与各分类的特征图叠加得到可视化图像。

此过程不需要重新训练模型,但要在模型中附加若干 hook,并进行若干手动的运算(如反向传播,加权求和),不易独自写出;因此,这次直接摘录他人已有代码中的相关方法,引入到自己的代码之中。在上面 6 张图像的目录下运行相应程序,得到的可视化结果如表 3 所示;篇幅起见,无论正误,仅展示每张图片第一预测结果的可视化图像。图片下方标注的为预测标签,加粗者表示预测正确。此外,由于这一步中的 6 张图片是压入一个 minibatch 预测的,故部分图片上的预测结果与之前有所不同(如之前预测错误,现在也一样的 2、3 号图片)。

表 3. Grad-CAM 可视化结果

手风琴 「烛台」 「脸」
小龙虾 「温莎椅」 阴阳

从上表可以发现,模型在小龙虾对应的 4 号图片中提取到了明显的特征,6 号图片中提取的特征也勉强可见(似乎是提取到了黑白色块的分界线,和圆周围的空白),而其他图片中的可视化结果则令人费解。特别是 1 号图片中,即使预测结果完全正确,也很难看出其是如何将这张图片识别为手风琴的。

仅就这 6 张图片来看,我们可以认为:尽管模型在一些分类上取得了正确的预测结果,但其是否真正提取到了合理的特征还有待考察。考虑到数据集本身容量很小,我们可以认为,数据太少是导致这种现象发生的原因,模型的泛化能力可能不足。

总结

在本次训练中,于数据准备方面做了相当多的编程练习(如划分数据集、翻译标签、估算 RGB 信息等),基本了解了迁移学习的操作流程,并巩固了上一次训练时初步尝试的陌生数据预测与数据可视化工作。

不过,从可视化结果中,可以看到模型还没有达到足够令人信服的程度,现在又回到之前没有能够解决的问题上来了:如何训练一个好的模型?好在这次掌握了更多的手段,对训练的过程认识更加清晰,现在在这一方面的准备也就更加充足。在之后的具体工作中,应该能够有所进步。