利用MMPreTrain微调图像分类模型

前言

• MMPreTrain是一款基于PyTorch的开源深度学习预工具箱，是OpenMMLab项目的成员之一
• MMPreTrain的主要特性有：
  • 支持多元化的主干网络与预训练模型
  • 支持多种训练策略（有监督学习，无监督学习，多模态学习等）
  • 提供多种训练技巧
  • 大量的训练配置文件
  • 高效率和高可扩展性
  • 功能强大的工具箱，有助于模型分析和实验
  • 支持多个开箱即用的推理任务
    • 图片分类
    • 图片描述（图片说明）
    • 视觉问答（Visual Question Answering）
    • 视觉定位（Visual Grounding）
    • 搜索（图搜图，图搜文，文搜图）
• 本文主要演示如何使用MMPreTrain，对图像分类任务，微调vision_transformer模型
• 分类数据使用kaggle平台中的Animal Faces数据集，运行环境为kaggle GPU P100

环境安装

• 因为需要对模型进行可解释分析，需要安装grad-cam包，mmcv的安装方式在我前面的mmdetection和mmsegmentation教程中都有写到。这里不再提

• mmpretrain安装方法最好是使用git，如果没有git工具，可以使用mim install mmpretrain

• 最后在项目文件夹下新建checkpoint、outputs、data文件夹，分别用来存放模型预训练权重、模型输出结果、训练数据

from IPython import display
!pip install "grad-cam>=1.3.6"
!pip install -U openmim
!pip install -q /kaggle/input/frozen-packages-mmdetection/mmcv-2.0.1-cp310-cp310-linux_x86_64.whl

!git clone https://github.com/open-mmlab/mmpretrain.git
%cd mmpretrain
!mim install -e .

!mkdir checkpoint
!mkdir outputs
!mkdir data

display.clear_output()

• 在上面的安装工作完成后，我们检查一下环境，以及核对一下安装版本

import mmcv
from mmcv.ops import get_compiling_cuda_version, get_compiler_version
import mmpretrain
print('MMCV版本', mmcv.__version__)
print('mmpretrain版本', mmpretrain.__version__)
print('CUDA版本', get_compiling_cuda_version())
print('编译器版本', get_compiler_version())

输出：

MMCV版本 2.0.1
mmpretrain版本 1.0.0
CUDA版本 11.8
编译器版本 GCC 11.3

• 因为mmpretrain适用于很多种任务，因此在进行图像分类模型微调时先查看一下支持该任务的模型有哪些，可以调用list_models函数查看，因为我们想要微调vision_transformer模型，可以加上限制条件vit进行更精准的筛选

from mmpretrain import list_models, inference_model
list_models(task='Image Classification',pattern='vit')

• 这里以候选列表中的vit-base-p32_in21k-pre_3rdparty_in1k-384px模型为例

模型推理

• 进入项目文件夹configs/vision_transformer，查看模型型号对应预训练权重及config文件

• 下载预训练权重，对示例图片进行推理
  

from mmpretrain import ImageClassificationInferencer
# 待输入图像路径
img_path = 'mmpretrain/demo/bird.JPEG'
model = 'vit-base-p32_in21k-pre_3rdparty_in1k-384px'
# 预训练权重
pretrained = './checkpoints/vit-base-p32_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-9cea8599.pth'
# 推理预测
inferencer = ImageClassificationInferencer(model=model, pretrained=pretrained, device='cuda:0')

result = inferencer('demo/bird.JPEG', show_dir="./visualize/")

display.clear_output()

• 查看推理结果

result[0].keys()

输出：

dict_keys(['pred_scores', 'pred_label', 'pred_score', 'pred_class'])

• 打印分类置信度最高类别的名称，以及置信度

# 置信度最高类别的名称
print(result[0]['pred_class'])
# 置信度最高类别的置信度
print('{:.3f}'.format(result[0]['pred_score']))

house finch, linnet, Carpodacus mexicanus
0.985

微调模型

• 将数据集移到data目录下，准备训练

# animal数据集移动
shutil.copytree('/kaggle/input/animal-faces/afhq', './data/animal')

配置文件解析

• MMPreTrain配置文件和mmdetection、mmsegmentation有点不太一样，当你打开vit-base-p32_in21k-pre_3rdparty_in1k-384px的配置文件vit-base-p32_64xb64_in1k-384px.py时，会发现配置文件中只显式的定义了数据管道及处理方式
• 但实际上数据处理和优化器参数都被隐形定义在_base_下了，详细情况可以看下面的代码注释

_base_ = [
    '../_base_/models/vit-base-p32.py',					# 模型配置
    '../_base_/datasets/imagenet_bs64_pil_resize.py',	# 数据配置
    '../_base_/schedules/imagenet_bs4096_AdamW.py',		# 训练策略配置
    '../_base_/default_runtime.py'						# 默认运行设置
]

# model setting
# 输入图像大小
model = dict(backbone=dict(img_size=384))  

# dataset setting
# 输入的图片数据通道以 'RGB' 顺序
data_preprocessor = dict(
    mean=[127.5, 127.5, 127.5],		# 输入图像归一化的 RGB 通道均值
    std=[127.5, 127.5, 127.5],		# 输入图像归一化的 RGB 通道标准差
    to_rgb=True,					# 是否将通道翻转，从 BGR 转为 RGB 或者 RGB 转为 BGR
)

train_pipeline = [
    dict(type='LoadImageFromFile'),			# 读取图像
    dict(type='RandomResizedCrop', scale=384, backend='pillow'),	# 随机放缩裁剪
    dict(type='RandomFlip', prob=0.5, direction='horizontal'),		# 随机水平翻转
    dict(type='PackInputs'),	# 准备图像以及标签
]

test_pipeline = [
    dict(type='LoadImageFromFile'),			# 读取图像
    dict(type='ResizeEdge', scale=384, edge='short', backend='pillow'),	# 缩放短边尺寸至384px
    dict(type='CenterCrop', crop_size=384),		# 中心裁剪
    dict(type='PackInputs'),		# 准备图像以及标签
]

train_dataloader = dict(dataset=dict(pipeline=train_pipeline))
val_dataloader = dict(dataset=dict(pipeline=test_pipeline))
test_dataloader = dict(dataset=dict(pipeline=test_pipeline))

# schedule设定
optim_wrapper = dict(clip_grad=dict(max_norm=1.0))

• 打开../_base_/models/vit-base-p32.py文件，查看模型配置

model = dict(
    type='ImageClassifier',		# 主模型类型（对于图像分类任务，使用 `ImageClassifier`）
    backbone=dict(
        type='VisionTransformer',	# 主干网络类型
        arch='b',
        img_size=224,		# 输入模型图像大小
        patch_size=32,		# patch数
        drop_rate=0.1,		# dropout率
        init_cfg=[			# 初始化参数方式
            dict(
                type='Kaiming',
                layer='Conv2d',
                mode='fan_in',
                nonlinearity='linear')
        ]),
    neck=None,
    head=dict(
        type='VisionTransformerClsHead',		# 分类颈网络类型
        num_classes=1000,						# 分类数
        in_channels=768,						# 输入通道数
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),	# 损失函数配置信息
        topk=(1, 5),	# 评估指标，Top-k 准确率， 这里为 top1 与 top5 准确率
    ))

• 打开../_base_/datasets/imagenet_bs64_pil_resize.py文件，查看数据配置

dataset_type = 'ImageNet'		# 预处理配置
data_preprocessor = dict(
    num_classes=1000,
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    to_rgb=True,
)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', scale=224, backend='pillow'),
    dict(type='RandomFlip', prob=0.5, direction='horizontal'),
    dict(type='PackInputs'),
]

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='ResizeEdge', scale=256, edge='short', backend='pillow'),
    dict(type='CenterCrop', crop_size=224),
    dict(type='PackInputs'),
]

train_dataloader = dict(
    batch_size=64,		# 每张 GPU 的 batchsize
    num_workers=5,		# 每个 GPU 的线程数
    dataset=dict(		# 训练数据集
        type=dataset_type,
        data_root='data/imagenet',
        split='train',
        pipeline=train_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=True),	# 默认采样器
)

val_dataloader = dict(
    batch_size=64,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='data/imagenet',
        split='val',
        pipeline=test_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=False),
)

# 验证集评估设置，使用准确率为指标， 这里使用 topk1 以及 top5 准确率
val_evaluator = dict(type='Accuracy', topk=(1, 5))

test_dataloader = val_dataloader
test_evaluator = val_evaluator

• 打开../_base_/schedules/imagenet_bs4096_AdamW.py文件，查看训练策略配置

optim_wrapper = dict(
    optimizer=dict(type='AdamW', lr=0.003, weight_decay=0.3),		# 使用AdamW优化器
    # vit预训练专用配置
    paramwise_cfg=dict(custom_keys={
        '.cls_token': dict(decay_mult=0.0),
        '.pos_embed': dict(decay_mult=0.0)
    }),
)

# 学习率策略
param_scheduler = [
    # 预热学习率调度器
    dict(
        type='LinearLR',
        start_factor=1e-4,
        by_epoch=True,
        begin=0,
        end=30,
        # 根据iter更新
        convert_to_iter_based=True),
    # 主要的学习策略
    dict(
        type='CosineAnnealingLR',
        T_max=270,
        by_epoch=True,
        begin=30,
        end=300,
    )
]

# train, val, test设置，max_epoch和验证频率
train_cfg = dict(by_epoch=True, max_epochs=300, val_interval=1)
val_cfg = dict()
test_cfg = dict()

auto_scale_lr = dict(base_batch_size=4096)

• 打开../_base_/default_runtime.py文件，查看默认运行设置

# 默认所有注册器使用的域
default_scope = 'mmpretrain'

# 配置默认的 hook
default_hooks = dict(
    # 记录每次迭代的时间
    timer=dict(type='IterTimerHook'),

    # 每 100 次迭代打印一次日志
    logger=dict(type='LoggerHook', interval=100),

    # 启用默认参数调度 hook
    param_scheduler=dict(type='ParamSchedulerHook'),

    # 每个 epoch 保存检查点
    checkpoint=dict(type='CheckpointHook', interval=1),

    # 在分布式环境中设置采样器种子
    sampler_seed=dict(type='DistSamplerSeedHook'),

    # 验证结果可视化，默认不启用，设置 True 时启用
    visualization=dict(type='VisualizationHook', enable=False),
)

# 配置环境
env_cfg = dict(
    # 是否开启 cudnn benchmark
    cudnn_benchmark=False,

    # 设置多进程参数
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),

    # 设置分布式参数
    dist_cfg=dict(backend='nccl'),
)

# 设置可视化工具
vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(type='UniversalVisualizer', vis_backends=vis_backends)

# 设置日志级别
log_level = 'INFO'

# 从哪个检查点加载
load_from = None

# 是否从加载的检查点恢复训练
resume = False

# 默认随机数种子
randomness = dict(seed=None, deterministic=False)

修改配置文件

• 根据上述说明，这里提供两种修改配置文件的方法。
• 第1种是将共5个配置文件的信息写在一个新的配置文件vit-base-p32_1xb64_in1k-384px_animal.py中，然后修改其中的内容。
• 首先将配置文件vit-base-p32_64xb64_in1k-384px.py中的内容更新到继承的键值对中，比如model中的img_size=384，train_pipeline和test_pipeline也都需要改
• 然后更改num_classes、dataset_type、train_dataloader、val_dataloader、val_evaluator、lr、param_scheduler、default_hooks、randomness
• 需要注意的是dataset_type要改成'CustomDataset'，而'CustomDataset'中是没有split键的，所以要删掉train_dataloader、val_dataloader中的split键
• 因为分类数量较少，不足5类，所以将val_evaluator中的topk由(1, 5)改成5
• lr要与原batch的lr进行等比例缩放，缩放率为32/(64 * 64)（由配置文件名64xb64可知，原batch_size为64 * 64）
• 因为只训练100个epoch，所以LinearLR scheduler中的end键也进行等比例缩放，即除以3。则CosineAnnealingLR scheduler中的T_max、begin、end相应变化
• 因为模型在前20个epoch可能没有学习成果，所以没有验证的必要，这里加入val_begin键，表示从第20个epoch开始在验证集上计算指标，验证的频率也不需要1个epoch一次，这里改成5个epoch验证一次
• 我们想要模型没10个epoch自动保存一次权重，并且最多同时保留两个训练权重，同时根据指标自动保留精度最高的训练权重checkpoint = dict(type='CheckpointHook', interval=10, max_keep_ckpts=2, save_best='auto')
• 记录频率100（单位：iter）有点太低，我们改成10
• 最后固定随机数种子randomness

custom_config = """
model = dict(
    type='ImageClassifier',		# 主模型类型（对于图像分类任务，使用 `ImageClassifier`）
    backbone=dict(
        type='VisionTransformer',	# 主干网络类型
        arch='b',
        img_size=384,		# 输入模型图像大小
        patch_size=32,		# patch数
        drop_rate=0.1,		# dropout率
        init_cfg=[			# 初始化参数方式
            dict(
                type='Kaiming',
                layer='Conv2d',
                mode='fan_in',
                nonlinearity='linear')
        ]),
    neck=None,
    head=dict(
        type='VisionTransformerClsHead',		# 分类颈网络类型
        num_classes=3,						# 分类数
        in_channels=768,						# 输入通道数
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),	# 损失函数配置信息
        topk=(1, 5),	# 评估指标，Top-k 准确率， 这里为 top1 与 top5 准确率
    ))

dataset_type = 'CustomDataset'		# 预处理配置
data_preprocessor = dict(
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    to_rgb=True,
)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', scale=384, backend='pillow'),
    dict(type='RandomFlip', prob=0.5, direction='horizontal'),
    dict(type='PackInputs'),
]

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='ResizeEdge', scale=384, edge='short', backend='pillow'),
    dict(type='CenterCrop', crop_size=384),
    dict(type='PackInputs'),
]

train_dataloader = dict(
    batch_size=64,		# 每张 GPU 的 batchsize
    num_workers=2,		# 每个 GPU 的线程数
    dataset=dict(		# 训练数据集
        type=dataset_type,
        data_root='./data/animal/train',
        pipeline=train_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=True),	# 默认采样器
)

val_dataloader = dict(
    batch_size=64,
    num_workers=2,
    dataset=dict(
        type=dataset_type,
        data_root='./data/animal/val',
        pipeline=test_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=False),
)

# 验证集评估设置，使用准确率为指标， 这里使用 topk1 以及 top5 准确率
val_evaluator = dict(type='Accuracy', topk=1)

test_dataloader = val_dataloader
test_evaluator = val_evaluator

optim_wrapper = dict(
    optimizer=dict(type='AdamW', lr=0.003 * 32 / (64 * 64), weight_decay=0.3),
    # vit预训练专用配置
    paramwise_cfg=dict(custom_keys={
        '.cls_token': dict(decay_mult=0.0),
        '.pos_embed': dict(decay_mult=0.0)
    }),
    clip_grad=dict(max_norm=1.0)
)

# 学习率策略
param_scheduler = [
    # 预热学习率调度器
    dict(
        type='LinearLR',
        start_factor=1e-4,
        by_epoch=True,
        begin=0,
        end=10,
        # 根据iter更新
        convert_to_iter_based=True),
    # 主要的学习策略
    dict(
        type='CosineAnnealingLR',
        T_max=90,
        by_epoch=True,
        begin=10,
        end=100,
    )
]

# train, val, test设置，max_epoch和验证频率
train_cfg = dict(by_epoch=True, max_epochs=100, val_begin=20, val_interval=5)
val_cfg = dict()
test_cfg = dict()

# 默认所有注册器使用的域
default_scope = 'mmpretrain'

# 配置默认的 hook
default_hooks = dict(
    # 记录每次迭代的时间
    timer=dict(type='IterTimerHook'),

    # 每 10 次迭代打印一次日志
    logger=dict(type='LoggerHook', interval=10),

    # 启用默认参数调度 hook
    param_scheduler=dict(type='ParamSchedulerHook'),

    checkpoint=dict(type='CheckpointHook', interval=10, max_keep_ckpts=2, save_best='auto'),

    # 在分布式环境中设置采样器种子
    sampler_seed=dict(type='DistSamplerSeedHook'),

    # 验证结果可视化，默认不启用，设置 True 时启用
    visualization=dict(type='VisualizationHook', enable=False),
)

# 配置环境
env_cfg = dict(
    # 是否开启 cudnn benchmark
    cudnn_benchmark=False,

    # 设置多进程参数
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),

    # 设置分布式参数
    dist_cfg=dict(backend='nccl'),
)

# 设置可视化工具
vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(type='UniversalVisualizer', vis_backends=vis_backends)

# 设置日志级别
log_level = 'INFO'

# 从哪个检查点加载
load_from = None

# 是否从加载的检查点恢复训练
resume = False

# 默认随机数种子
randomness = dict(seed=2023, deterministic=False)
"""
# 写入vit-base-p32_1xb64_in1k-384px_pets.py文件中
animal_config=f'./configs/vision_transformer/vit-base-p32_1xb64_in1k-384px_pets.py'
with open(animal_config, 'w') as f:
    f.write(custom_config)

• 第2种方法是先将默认配置文件读取，然后再通过python的字典特性进行更改，优点是，只改动需要改动的地方，逻辑清晰
• 缺点是有些在配置文件中的中间变量无效了，比如在配置文件中可以只用定义dataset_type，后面train_dataloader、val_dataloader可以直接使用，但是用字典特性要改两遍
• 参数的改变和上面的一样，但是代码少很多

# 读取配置文件
from mmengine import Config
cfg = Config.fromfile('./configs/vision_transformer/vit-base-p32_64xb64_in1k-384px.py')

max_epochs = 100
batch_size = 64
lr_scale_factor = batch_size/(64 * 64)
epoch_scale_factor = max_epochs/cfg.train_cfg.max_epochs

cfg.model.head.num_classes = 3

cfg.load_from = './checkpoints/vit-base-p32_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-9cea8599.pth'
cfg.work_dir = './work_dir'

cfg.dataset_type = 'CustomDataset'

cfg.train_dataloader.batch_size = batch_size
cfg.train_dataloader.num_workers = 2
cfg.train_dataloader.dataset.type = cfg.dataset_type
cfg.train_dataloader.dataset.data_root = './data/animal/train'
del cfg.train_dataloader.dataset['split']

cfg.val_dataloader.batch_size = cfg.train_dataloader.batch_size
cfg.val_dataloader.num_workers = cfg.train_dataloader.num_workers
cfg.val_dataloader.dataset.data_root = './data/animal/valid'
cfg.val_dataloader.dataset.type = cfg.dataset_type
del cfg.val_dataloader.dataset['split']

cfg.test_dataloader = cfg.val_dataloader

cfg.val_evaluator = dict(type='Accuracy', topk=1)
cfg.test_evaluator = cfg.val_evaluator

cfg.optim_wrapper.optimizer.lr = cfg.optim_wrapper.optimizer.lr * lr_scale_factor

# LinearLR scheduler end epoch
cfg.param_scheduler[0].end = cfg.param_scheduler[0].end * epoch_scale_factor

# CosineAnnealingLR scheduler
cfg.param_scheduler[1].T_max = max_epochs - cfg.param_scheduler[0].end
cfg.param_scheduler[1].begin = cfg.param_scheduler[0].end
cfg.param_scheduler[1].end = max_epochs

cfg.train_cfg.max_epochs = max_epochs
cfg.train_cfg.val_begin = 20
cfg.train_cfg.val_interval = 5

cfg.default_hooks.checkpoint = dict(type='CheckpointHook', interval=10, max_keep_ckpts=2, save_best='auto')
cfg.default_hooks.logger.interval = 50

cfg.randomness.seed = 2023

#------------------------------------------------------
animal_config=f'./configs/vision_transformer/vit-base-p32_1xb64_in1k-384px_pets.py'
with open(animal_config, 'w') as f:
    f.write(cfg.pretty_text)

启动训练

!python tools/train.py {animal_config}

• 由于输出日志太长，这里就不全部展示了，打印一下精度最高的模型权重

07/30 13:33:50 - mmengine - INFO - Epoch(val) [55][24/24]    accuracy/top1: 99.9333  data_time: 0.2443  time: 0.5068

• 可以看到模型在验证集上的精度为99.93%，可以说非常不错

模型推理

• 加载精度最高的模型，并对图片进行推理

import glob
ckpt_path = glob.glob('./work_dir/best_accuracy_top1*.pth')[0]
img_path = '/kaggle/input/animal-faces/afhq/train/cat/flickr_cat_000052.jpg'

inferencer = ImageClassificationInferencer(animal_config, pretrained=ckpt_path)

result = inferencer(img_path)
result

输出：

[{'pred_scores': array([9.9998045e-01, 1.3512783e-05, 6.0256166e-06], dtype=float32),
  'pred_label': 0,
  'pred_score': 0.9999804496765137,
  'pred_class': 'cat'}]

绘制混淆矩阵

• 我们可以绘制混淆矩阵进一步检查模型精度

python tools/analysis_tools/confusion_matrix.py \
    {animal_config} \
    {ckpt_path}\
    --show

类别激活图（CAM）可视化

• 使用类别激活图（CAM）对分类图像进行解释，更多参数设置请参照官方文档

!python tools/visualization/vis_cam.py \
        {img_path} \
        {animal_config} \
        {ckpt_path} \
        --method GradCAM \
        --save-path cam.jpg \
        --vit-like
display.clear_output()

from PIL import Image
Image.open('cam.jpg')

T-SNE可视化

• 通过降维可视化可以进一步观察模型对类别分界线是否明显，还可以找到模型容易误判的类别

python tools/visualization/vis_tsne.py \
     {animal_config}\
    --checkpoint {ckpt_path}