利用MMPreTrain微调图像分类模型

MMPreTrain图像分类任务模型微调教程

最新推荐文章于 2025-09-22 14:52:51 发布

原创

最新推荐文章于 2025-09-22 14:52:51 发布 · 7.5k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#分类 #人工智能 #MMPreTrain

本文介绍了基于MMPreTrain的开源深度学习预训练工具箱，其支持多元化主干网络与预训练模型等。主要演示使用MMPreTrain对图像分类任务微调模型，包括环境安装、模型推理、微调模型（配置文件解析、修改、启动训练），还进行了绘制混淆矩阵、CAM可视化、T - SNE可视化等操作。

前言

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

微调模型

# 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'