Skip to content

FOMO 模型训练

本节将介绍如何在 COCO MASK 数据集上训练 FOMO 口罩检测模型。FOMO 口罩检测模型基于 MobileNet V2 和 MobileNet V3 实现 (实际选用的神经网络取决于您选择的模型配置文件)。

关于 MobileNet 的更多信息,请参考论文 MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications

准备数据集

SSCMA 默认使用 COCO_MASK 数据集训练 FOMO 模型,请参照以下步骤完成数据集的准备。

  1. 参考互联网数据集 - SSCMA - COCO_MASK 数据集下载数据集并完成数据集的解压。

  2. 记住数据集解压后的文件夹路径 (如: datasets\mask),在之后修改配置文件时需要使用该文件夹路径。

选择配置文件

我们将根据需要执行的训练任务类型来选择合适的配置文件,我们已经在模型配置中对配置文件的功能、结构、原理进行了简单介绍。

对于 FOMO 模型示例,我们使用 fomo_mobnetv2_0.35_x8_abl_coco.py 作为配置文件,它位于 SSCMA 主目录路径 configs/fomo 下的文件夹中,并额外继承了 default_runtime_det.py 配置文件。

配置文件内容如下,对于初学者,我们建议首先注意该配置文件中 data_rootepochs 这两个参数。

fomo_mobnetv2_0.35_x8_abl_coco.py
python
_base_='../_base_/default_runtime_det.py'
default_scope='sscma'
custom_imports=dict(imports=['sscma'], allow_failed_imports=False)

num_classes=2
model=dict(type='Fomo',
           backbone=dict(type='mmdet.MobileNetV2', widen_factor=0.35, out_indices=(2,)),
           head=dict(type='FomoHead',
                     input_channels=[16],
                     num_classes=num_classes,
                     middle_channel=48,
                     act_cfg='ReLU6',
                     loss_cls=dict(type='BCEWithLogitsLoss',
                                   reduction='none',
                                   pos_weight=40),
                     loss_bg=dict(type='BCEWithLogitsLoss', reduction='none'),
           ),
)

# dataset settings
dataset_type='FomoDatasets'
data_root=''
height=96
width=96
batch_size=16
workers=1

train_pipeline=[
    dict(type='RandomResizedCrop',
         height=height,
         width=width,
         scale=(0.80, 1.2),
         p=1),
    dict(type='Rotate', limit=30),
    dict(type='RandomBrightnessContrast',
         brightness_limit=0.3,
         contrast_limit=0.3,
         p=0.5),
    dict(type='HorizontalFlip', p=0.5),
]
test_pipeline=[dict(type='Resize', height=height, width=width, p=1)]

train_dataloader=dict(
    batch_size=batch_size,
    num_workers=workers,
    persistent_workers=True,
    drop_last=False,
    collate_fn=dict(type='fomo_collate'),
    sampler=dict(type='DefaultSampler', shuffle=True, round_up=False),
    dataset=dict(type=dataset_type,
                 data_root=data_root,
                 ann_file='train/_annotations.coco.json',
                 img_prefix='train',
                 pipeline=train_pipeline),
)
val_dataloader=dict(
    batch_size=1,
    num_workers=1,
    persistent_workers=True,
    drop_last=False,
    collate_fn=dict(type='fomo_collate'),
    sampler=dict(type='DefaultSampler', shuffle=True, round_up=False),
    dataset=dict(type=dataset_type,
                 data_root=data_root,
                 ann_file='valid/_annotations.coco.json',
                 img_prefix='valid',
                 pipeline=test_pipeline))
test_dataloader=val_dataloader

# optimizer
lr=0.001
epochs=300
find_unused_parameters=True
optim_wrapper=dict(optimizer=dict(type='Adam', lr=lr, weight_decay=5e-4,eps=1e-7))

#evaluator
val_evaluator=dict(type='FomoMetric')
test_evaluator=val_evaluator
train_cfg=dict(by_epoch=True,max_epochs=70)

# learning policy
param_scheduler=[
    dict(type='LinearLR', begin=0, end=30, start_factor=0.001, by_epoch=False),  # warm-up
    dict(type='MultiStepLR',
         begin=1,
         end=500,
         milestones=[100, 200, 250],
         gamma=0.1,
         by_epoch=True)
]

训练模型

训练模型需要使用我们之前配置好的 SSCMA 工作环境,如果您按照我们的安装指南使用 Conda 将 SSCMA 安装在了名为 sscma 的虚拟环境中,请首先确保您目前正处在虚拟环境中。

然后,在 SSCMA 项目根目录,我们执行如下的指令,训练一个 FOMO 口罩识别模型。

sh
python3 tools/train.py \
    configs/fomo/fomo_mobnetv2_0.35_x8_abl_coco.py \
    --cfg-options \
        data_root='datasets/mask' \
        epochs=50

在训练期间,训练得到的模型权重和相关的日志信息会默认保存至路径 work_dirs/fomo_mobnetv2_0.35_x8_abl_coco 下,您可以使用 TensorBoard 等工具事实监测训练情况。

sh
tensorboard --logdir work_dirs/fomo_mobnetv2_0.35_x8_abl_coco

在训练完成后,最新的 FOMO 模型权重文件的路径会保存在 work_dirs/fomo_mobnetv2_0.35_x8_abl_coco/last_checkpoint 文件中。请留意权重文件路径,在将模型转换为其他格式时需要用到。

TIP

如果您配置了虚拟环境但并未激活,您可以使用以下命令激活虚拟环境:

sh
conda activate sscma

测试和评估

测试

在完成了 FOMO 模型的训练后,您可以使用以下命令,指定特定权重并测试模型:

sh
python3 tools/inference.py \
    configs/fomo/fomo_mobnetv2_0.35_x8_abl_coco.py \
    "$(cat work_dirs/fomo_mobnetv2_0.35_x8_abl_coco/last_checkpoint)" \
    --cfg-options \
        data_root='datasets/mask'

TIP

如果您在测试时希望实时的可视化预览,您可以在测试命令后追加一个参数 --show 来显示预测结果。对于更多可选的参数,请参考源代码 tools/test.py

评估

为了进一步在现实中的边缘计算设备上测试和评估模型,您需要导出模型。在导出模型的过程中,SSCMA 会对模型进行一些优化,如模型的剪枝、蒸馏等。您可以参考模型导出章节进一步学习如何导出模型。

部署

在导出模型后,你可以将模型部署到边缘计算设备上进行测试和评估。你可以参考 Deploy 部分来了解更多关于如何部署模型的信息。

在 Apache 2.0 授权下发布