一、MMClassification工具包简介
MMClassification工具包目前集成了包括VGG、ResNet、ConvNext以及Swin-Transformer等多种神经网络模型,同时为了便于使用,提供了包括MNIST、CIFAR以及ImageNet 1K和21K等多种图像分类数据集支持,其他数据集需求可在OpenDataLab中进行获取。
训练与推理使用
训练
针对训练任务,在通过源码安装好框架后,采用如下命令即可进行训练
# 单机单卡
python tools/train.py ${CONFIG_FILE} [optional arguments]
# 单机、多机多卡
./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments]
# 使用任务调度器 Slurm 启动任务
[GPUS=${GPUS}] ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${WORK_DIR}
# 加载预训练权重,从 checkpoint 恢复训练
上述命令后增加 --resume-from ${CHECKPOINT_FILE} 参数
推理
使用推理工具时同样需要源码安装,随后通过下述命令进行推理
# 单张图像
python demo/image_demo.py ${IMAGE_FILE} ${CONFIG_FILE} ${CHECKPOINT_FILE}
# 测试集
## 单卡
python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--metrics ${METRICS}] [--out
${RESULT_FILE}]
## 多卡
./tools/dist_test.sh ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} [--metrics ${METRICS}] [--out
${RESULT_FILE}]
推理API简介
在进行推理前,若没有自己的训练权重,需要通过OpenMMLab提供的mim管理工具下载对应配置文件和训练权重
$ mim download mmcls --config mobilenet-v2_8xb32_in1k --dest .
以MobileNet-v2为例,在进行推理时,仅需关注下述代码即可
from mmcls.apis import init_model, inference_model, show_result_pyplot
model = init_model('mobilenet-v2_8xb32_in1k.py',
'mobilenet_v2_batch256_imagenet_20200708-3b2dc3af.pth',
device='cuda:0')
result = inference_model(model, 'banana.png')
show_result_pyplot(model, 'banana.png', result)
- init_model():负责调用下载好的config文件、训练权重以及指定支持设备
- inference_model():推理API接口,负责返回推理结果
- show_result_pyplot():将推理结果展示在图像中
配置文件
在深度学习任务中,通常一个模型想正常训练,需要具备以下几个方面
- 模型结构:定义模型层数、每层通道数等具体结构
- 数据集:包括数据集划分策略、数据文件路径以及数据增强策略等
- 训练策略:决定模型在训练时采用的优化方法、学习率参数、batch_size、epoch以及学习率变化策略等
- 运行:GPU支持、分布式环境支持等
- 辅助功能:log打印、定时保存checkpoint以及可视化loss曲线
OpenMMLab将上述内容统一集成在了一个配置文件中,由多个字段仅需定义,用户仅需通过定义和修改该配置文件,即可定义自己的网络结构、数据内容以及训练策略,从而轻松完成整个训练过程。
- model:定义模型
- data:定义数据
- optimizer、lr_config:定义训练策略
- load_from:定义预训练模型的参数文件