mmyolo标注+训练+测试+部署全流程

官方教学视频：玩转 MMYOLO 之实用篇（三）：自定义数据集从标注到部署保姆级教程_哔哩哔哩_bilibili

从用户自定义图片数据集标注到最终进行训练和部署的整体流程。步骤概览如下：

1. 数据集准备：tools/misc/download_dataset.py

2. 使用 labelme 和算法进行辅助和优化数据集标注：demo/image_demo.py + labelme

3. 使用脚本转换成 COCO 数据集格式：tools/dataset_converters/labelme2coco.py

4. 数据集划分为训练集、验证集和测试集：tools/misc/coco_split.py

5. 根据数据集内容新建 config 文件

6. 数据集可视化分析：tools/analysis_tools/dataset_analysis.py

7. 优化 Anchor 尺寸：tools/analysis_tools/optimize_anchors.py

8. 可视化 config 配置中数据处理部分： tools/analysis_tools/browse_dataset.py

9. 训练：tools/train.py

10. 推理：demo/image_demo.py

11. 部署

1. 数据集准备

该数据集形式可作为参考，可以直接训练

└── ./data/cat
    ├── images # 图片文件
    │    ├── image1.jpg
    │    ├── image2.png
    │    └── ...
    ├── labels # labelme 标注文件
    │    ├── image1.json
    │    ├── image2.json
    │    └── ...
    ├── annotations # 数据集划分的 COCO 文件
    │    ├── annotations_all.json # 全量数据的 COCO label 文件
    │    ├── trainval.json # 划分比例 80% 的数据
    │    └── test.json # 划分比例 20% 的数据
    └── class_with_id.txt # id + class_name 文件

已经有数据，可以将其组成下面的结构：.
└── $DATA_ROOT
    └── images
         ├── image1.jpg
         ├── image2.png
         └── ...

2. 使用 labelme 和算法进行辅助和优化数据集标注

目前只使用过人工标注：

• 安装 labelme

  conda create -n labelme python=3.8
  conda activate labelme
  pip install labelme==5.1.1

• 启动 labelme

  labelme ${图片文件夹路径（即上一步的图片文件夹）} \
          --output ${label文件所处的文件夹路径（即上一步的 --out-dir）} \
          --autosave \
          --nodata

  其中：

• --output：labelme 标注文件保存路径，如果该路径下已经存在部分图片的标注文件，则会进行加载；

• --autosave：标注文件自动保存，会略去一些繁琐的保存步骤；

• --nodata：每张图片的标注文件中不保存图片的 base64 编码，设置了这个 flag 会大大减少标注文件的大小。

• 例子： 

  cd /path/to/mmyolo
  labelme ./data/cat/images --output ./data/cat/labels --autosave --nodata

  

个人常用方法：在lableme目录下直接运行labelImg.py

python labelImg.py

opendir打开需要标注的图像文件夹

save指定标注后保存的文件夹

3. 使用脚本转换成 COCO 数据集格式

3.1 MMYOLO 提供脚本将 labelme 的 label 转换为 COCO label

python tools/dataset_converters/labelme2coco.py --img-dir ${图片文件夹路径} \
                                                --labels-dir ${label 文件夹位置} \
                                                --out ${输出 COCO label json 路径} \
                                                [--class-id-txt ${class_with_id.txt 路径}]

其中： --class-id-txt：是数据集 id class_name 的 .txt 文件：

• 如果不指定，则脚本会自动生成，生成在 --out 同级的目录中，保存文件名为 class_with_id.txt；

• 如果指定，脚本仅会进行读取但不会新增或者覆盖，同时，脚本里面还会判断是否存在 .txt 中其他的类，如果出现了会报错提示，届时，请用户检查 .txt 文件并加入新的类及其 id。

例子：

以本教程的 cat 数据集为例：

python tools/dataset_converters/labelme2coco.py --img-dir ./data/cat/images \
                                                --labels-dir ./data/cat/labels \
                                                --out ./data/cat/annotations/annotations_all.json

 脚本：

import xml.etree.ElementTree as ET
import os
import json

coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

category_set = dict()
image_set = set()

# category_item_id = -1
category_item_id = 0
image_id = 0
annotation_id = 0


def addCatItem(name):
    global category_item_id
    category_item = dict()
    category_item['supercategory'] = 'none'
    category_item_id += 1
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_set[name] = category_item_id
    return category_item_id


def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    print(file_name)
    image_item['file_name'] = file_name + ".jpg"
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id


def addAnnoItem(object_name, image_id, category_id, bbox):
    global annotation_id
    annotation_item = dict()
    annotation_item['segmentation'] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])

    annotation_item['segmentation'].append(seg)

    annotation_item['area'] = bbox[2] * bbox[3]
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)


def parseXmlFiles(xml_path):
    for f in os.listdir(xml_path):
        if not f.endswith('.xml'):
            continue
        xmlname = f.split('.xml')[0]

        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None

        xml_file = os.path.join(xml_path, f)
        print(xml_file)

        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))

        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None

            if elem.tag == 'folder':
                continue

            if elem.tag == 'filename':
                file_name = xmlname
                if file_name in category_set:
                    raise Exception('file_name duplicated')

            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:

                    raise Exception('duplicated image: {}'.format(file_name))

                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>
            for subelem in elem:
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None

                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name':
                    object_name = subelem.text
                    if object_name not in category_set:
                        current_category_id = addCatItem(object_name)
                    else:
                        current_category_id = category_set[object_name]

                elif current_parent == 'size':
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text)

                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox':
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = int(float(option.text))

                # only after parse the <object> tag
                if bndbox['xmin'] is not None:
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['xmin'])
                    # y
                    bbox.append(bndbox['ymin'])
                    # w
                    bbox.append(bndbox['xmax'] - bndbox['xmin'])
                    # h
                    bbox.append(bndbox['ymax'] - bndbox['ymin'])
                    print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox)


if __name__ == '__main__':
    # 主要修改存放xml和生成json路径
    xml_path = '/home/nvidia/xxxx/xxxx/train_data/imageA/labels' 
    json_file ='/home/nvidia/xxxx/xxx/train_data/imageA/annotations/annotations_all.json'
    parseXmlFiles(xml_path)
    json.dump(coco, open(json_file, 'w'))

3.2 检查转换的 COCO label（从这开始需要在安装好的mmyolo文件夹下运行）

使用下面的命令可以将 COCO 的 label 在图片上进行显示，这一步可以验证刚刚转换是否有问题：

python tools/analysis_tools/browse_coco_json.py --img-dir ${图片文件夹路径} \
                                                --ann-file ${COCO label json 路径}

例子：

python tools/analysis_tools/browse_coco_json.py --img-dir ./data/cat/images \
                                                --ann-file ./data/cat/annotations/annotations_all.json

4. 数据集划分为训练集、验证集和测试集

通常，自定义图片都是一个大文件夹，里面全部都是图片，需要我们自己去对图片进行训练集、验证集、测试集的划分，如果数据量比较少，可以不划分验证集。下面是划分脚本的具体用法：

python tools/misc/coco_split.py --json ${COCO label json 路径} \
                                --out-dir ${划分 label json 保存根路径} \
                                --ratios ${划分比例} \
                                [--shuffle] \
                                [--seed ${划分的随机种子}]

其中：

• --ratios：划分的比例，如果只设置了 2 个，则划分为 trainval + test，如果设置为 3 个，则划分为 train + val + test。支持两种格式 —— 整数、小数：

  • 整数：按比例进行划分，代码中会进行归一化之后划分数据集。例子： --ratio 2 1 1（代码里面会转换成 0.5 0.25 0.25） or --ratio 3 1（代码里面会转换成 0.75 0.25）

  • 小数：划分为比例。如果加起来不为 1 ，则脚本会进行自动归一化修正。例子： --ratio 0.8 0.1 0.1 or --ratio 0.8 0.2

• --shuffle: 是否打乱数据集再进行划分；

• --seed：设定划分的随机种子，不设置的话自动生成随机种子。

例子：

python tools/misc/coco_split.py --json ./data/cat/annotations/annotations_all.json \
                                --out-dir ./data/cat/annotations \
                                --ratios 0.8 0.2 \
                                --shuffle \
                                --seed 10

5. 根据数据集内容新建 config 文件

确保数据集目录是这样的：

.
└── $DATA_ROOT
    ├── annotations
    │    ├── trainval.json # 根据上面的指令只划分 trainval + test，如果您使用 3 组划分比例的话，这里是 train.json、val.json、test.json
    │    └── test.json
    ├── images
    │    ├── image1.jpg
    │    ├── image1.png
    │    └── ...
    └── ...

因为是自定义的数据集，所以需要自己新建一个 config 并加入需要修改的部分信息。

关于新的 config 的命名：

• 这个 config 继承的是 yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py；

• 训练的类以本教程提供的数据集中的类 cat 为例（如果是自己的数据集，可以自定义类型的总称）；

• 本教程测试的显卡型号是 1 x 3080Ti 12G 显存，电脑内存 32G，可以训练 YOLOv5-s 最大批次是 batch size = 32（详细机器资料可见附录）；

• 训练轮次是 100 epoch。

综上所述：可以将其命名为 yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py，并将其放置在文件夹 configs/custom_dataset 中。

我们可以在 configs 目录下新建一个新的目录 custom_dataset，同时在里面新建该 config 文件，并添加以下内容：

_base_ = '../yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py'

max_epochs = 100  # 训练的最大 epoch
data_root = './data/cat/'  # 数据集目录的绝对路径
# data_root = '/root/workspace/mmyolo/data/cat/'  # Docker 容器里面数据集目录的绝对路径

# 结果保存的路径，可以省略，省略保存的文件名位于 work_dirs 下 config 同名的文件夹中
# 如果某个 config 只是修改了部分参数，修改这个变量就可以将新的训练文件保存到其他地方
work_dir = './work_dirs/yolov5_s-v61_syncbn_fast_1xb32-100e_cat'

# load_from 可以指定本地路径或者 URL，设置了 URL 会自动进行下载，因为上面已经下载过，我们这里设置本地路径
# 因为本教程是在 cat 数据集上微调，故这里需要使用 `load_from` 来加载 MMYOLO 中的预训练模型，这样可以在加快收敛速度的同时保证精度
load_from = './work_dirs/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth'  # noqa

# 根据自己的 GPU 情况，修改 batch size，YOLOv5-s 默认为 8卡 x 16bs
train_batch_size_per_gpu = 32
train_num_workers = 4  # 推荐使用 train_num_workers = nGPU x 4

save_epoch_intervals = 2  # 每 interval 轮迭代进行一次保存一次权重

# 根据自己的 GPU 情况，修改 base_lr，修改的比例是 base_lr_default * (your_bs / default_bs)
base_lr = _base_.base_lr / 4

anchors = [  # 此处已经根据数据集特点更新了 anchor，关于 anchor 的生成，后面小节会讲解
    [(68, 69), (154, 91), (143, 162)],  # P3/8
    [(242, 160), (189, 287), (391, 207)],  # P4/16
    [(353, 337), (539, 341), (443, 432)]  # P5/32
]

class_name = ('cat', )  # 根据 class_with_id.txt 类别信息，设置 class_name
num_classes = len(class_name)
metainfo = dict(
    classes=class_name,
    palette=[(220, 20, 60)]  # 画图时候的颜色，随便设置即可
)

train_cfg = dict(
    max_epochs=max_epochs,
    val_begin=20,  # 第几个 epoch 后验证，这里设置 20 是因为前 20 个 epoch 精度不高，测试意义不大，故跳过
    val_interval=save_epoch_intervals  # 每 val_interval 轮迭代进行一次测试评估
)

model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        prior_generator=dict(base_sizes=anchors),

        # loss_cls 会根据 num_classes 动态调整，但是 num_classes = 1 的时候，loss_cls 恒为 0
        loss_cls=dict(loss_weight=0.5 *
                      (num_classes / 80 * 3 / _base_.num_det_layers))))

train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        _delete_=True,
        type='RepeatDataset',
        # 数据量太少的话，可以使用 RepeatDataset ，在每个 epoch 内重复当前数据集 n 次，这里设置 5 是重复 5 次
        times=5,
        dataset=dict(
            type=_base_.dataset_type,
            data_root=data_root,
            metainfo=metainfo,
            ann_file='annotations/trainval.json',
            data_prefix=dict(img='images/'),
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            pipeline=_base_.train_pipeline)))

val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='annotations/trainval.json',
        data_prefix=dict(img='images/')))

test_dataloader = val_dataloader

val_evaluator = dict(ann_file=data_root + 'annotations/trainval.json')
test_evaluator = val_evaluator

optim_wrapper = dict(optimizer=dict(lr=base_lr))

default_hooks = dict(
    # 设置间隔多少个 epoch 保存模型，以及保存模型最多几个，`save_best` 是另外保存最佳模型（推荐）
    checkpoint=dict(
        type='CheckpointHook',
        interval=save_epoch_intervals,
        max_keep_ckpts=5,
        save_best='auto'),
    param_scheduler=dict(max_epochs=max_epochs),
    # logger 输出的间隔
    logger=dict(type='LoggerHook', interval=10))

实操过程：

# 验证数据集
# 在该路径下/mmyolo/tools/analysis_tools
python analysis_tools/browse_coco_json.py --img-dir /home/nvidia/xx/Projects/xx/xx/shiwai_fire/train/images --ann-file /home/nvidia/xx/Projects/xx/xx/shiwai_fire/train/annotations/annotations_all.json

# 划分数据集
# 在该路径下/mmyolo/tools/misc
python coco_split.py --json /home/nvidia/YQ/Projects/BaoGang/hq/shiwai_fire/train/annotations/annotations_all.json --out-dir /home/nvidia/YQ/Projects/BaoGang/hq/shiwai_fire/train/annotations --ratios 0.8 0.2 --shuffle --seed 10

# 训练模型
# 在该路径下/mmyolo
python tools/train.py configs/yolov5/yolov5_s-v61_fast_1xb12-40e_608x352_xxxx(根据自己需求更改).py

# 继续训练
# 在该路径下/mmyolo
python tools/train.py configs/yolov5/yolov5_s-v61_fast_1xb12-40e_608x352_xxxx.py --resume

 自己的模型代码：yolov5_s-v61_fast_1xb12-40e_608x352_xxx.py

_base_ = 'yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py'

data_root = '/home/zeng/project/train/'
# 训练的类型名，要与标注的相对应
class_name = ('xwq')
# 若是多个
# class_name = ('xwq', "xx")

# 训练的类型有多少个
num_classes = 1
metainfo = dict(classes=class_name, palette=[(20, 220, 60)])

anchors = [
    [(68, 69), (154, 91), (143, 162)],  # P3/8
    [(242, 160), (189, 287), (391, 207)],  # P4/16
    [(353, 337), (539, 341), (443, 432)]  # P5/32
]
# 根据实际图片分辨率进行更改
img_scale = (640, 480)
max_epochs = 200
train_batch_size_per_gpu = 14
train_num_workers = 2

load_from = 'https://download.openmmlab.com/mmyolo/v0/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth'  # noqa


model = dict(
    backbone=dict(frozen_stages=4),
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        prior_generator=dict(base_sizes=anchors)))

train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        # 训练数据的json存放路径
        ann_file='annotations/trainval.json',
        # 训练数据的图片存放路径
        data_prefix=dict(img='images/')))

val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        # 测试数据的json存放路径
        ann_file='annotations/test.json',
        # 测试数据的图片存放路径
        data_prefix=dict(img='images/')))

test_dataloader = val_dataloader

_base_.optim_wrapper.optimizer.batch_size_per_gpu = train_batch_size_per_gpu

val_evaluator = dict(ann_file=data_root + 'annotations/test.json')
test_evaluator = val_evaluator

default_hooks = dict(
    checkpoint=dict(interval=20, max_keep_ckpts=2, save_best='auto'),
    # The warmup_mim_iter parameter is critical.
    # The default value is 1000 which is not suitable for cat datasets.
    param_scheduler=dict(max_epochs=max_epochs, warmup_mim_iter=10),
    logger=dict(type='LoggerHook', interval=5))
train_cfg = dict(max_epochs=max_epochs, val_interval=10)
#visualizer = dict(vis_backends = [dict(type='LocalVisBackend'), dict(type='WandbVisBackend')]) # noqa
#visualizer = dict(vis_backends=[dict(type='LocalVisBackend'), dict(type='TensorboardVisBackend')])

训练完成后 利用test.py进行测试

from mmdet.apis import init_detector, inference_detector
# from mmdet.utils import register_all_modules
# from mmdet.registry import VISUALIZERS
# import mmcv
import random
import cv2
import os
import numpy as np
config_file = '/media/nvidia/HP P500/粮仓/yolov5_s-v61_fast_1xb12-40e_608x352_topdown.py'
checkpoint_file = '/media/nvidia/HP P500/粮仓/best_coco_bbox_mAP_epoch_60.pth'
model = init_detector(config_file, checkpoint_file, device='cuda:0')  # or device='cuda:0'

# 验证集数据路径
imgPath = "/media/nvidia/xxxx/xxx/yanzheng"
# imgPath = "/home/main/NewData/Preception/xxx/train_datas/datas/test_imgs/right"
imgLists = os.listdir(imgPath)
# saveImgPath = "/home/main/NewData/Preception/托盘/train_datas/datas/test_imgs/result_right"
# 结果保存路径
saveImgPath = "/media/nvidia/xxx/xxx/result"
# for imgName in imgLists[1:2]:
for imgName in imgLists:
    imgpath = os.path.join(imgPath, imgName)
    print(imgpath)
    # img2 = cv2.imread(imgpath)
    # result = inference_detector(model, img2)
    result = inference_detector(model, imgpath)
    # print(result)
    # pred_instances = result.pred_instances[
    #             result.pred_instances.scores > 0.80]
    # print("pred_instances", pred_instances)
    img = cv2.imread(imgpath)
    position = (50,50)
    font = cv2.LINE_AA
    font_scale = 2
    font_thickness = 5
    font_color = (0, 0, 255)  # 红色
    text = "score is: {}".format(result.pred_instances.scores[0].item())
    cv2.putText(img, text, position, font, font_scale, font_color, font_thickness)
    for result in result.pred_instances:
        # 分数达到多少则进行显示
        if result.scores.item() > 0.4:

    # if result.pred_instances.scores[0].item() > 0.60:

            x1 = int(result.bboxes[0][0].item())
            y1 = int(result.bboxes[0][1].item())
            x2 = int(result.bboxes[0][2].item())
            y2 = int(result.bboxes[0][3].item())
            
            cv2.rectangle(img, (x1,y1), (x2,y2), (0, 255, 255), 3)
            cv2.imwrite(os.path.join(saveImgPath, imgName), img)
        

通过验证集可以得到初步训练完模型的效果如何，若满意，则进行部署

 6. 尝试 MMYOLO 其他模型

MMYOLO 集成了多种 YOLO 算法，切换非常方便，无需重新熟悉一个新的 repo，直接切换 config 文件就可以轻松切换 YOLO 模型，只需简单 3 步即可切换模型：

1. 新建 config 文件

2. 下载预训练权重

3. 启动训练

下面以 YOLOv6-s 为例，进行讲解。

1. 搭建一个新的 config：

   _base_ = '../yolov6/yolov6_s_syncbn_fast_8xb32-400e_coco.py'
   
   max_epochs = 100  # 训练的最大 epoch
   data_root = './data/cat/'  # 数据集目录的绝对路径
   
   # 结果保存的路径，可以省略，省略保存的文件名位于 work_dirs 下 config 同名的文件夹中
   # 如果某个 config 只是修改了部分参数，修改这个变量就可以将新的训练文件保存到其他地方
   work_dir = './work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat'
   
   # load_from 可以指定本地路径或者 URL，设置了 URL 会自动进行下载，因为上面已经下载过，我们这里设置本地路径
   # 因为本教程是在 cat 数据集上微调，故这里需要使用 `load_from` 来加载 MMYOLO 中的预训练模型，这样可以在加快收敛速度的同时保证精度
   load_from = './work_dirs/yolov6_s_syncbn_fast_8xb32-400e_coco_20221102_203035-932e1d91.pth'  # noqa
   
   # 根据自己的 GPU 情况，修改 batch size，YOLOv6-s 默认为 8卡 x 32bs
   train_batch_size_per_gpu = 32
   train_num_workers = 4  # 推荐使用 train_num_workers = nGPU x 4
   
   save_epoch_intervals = 2  # 每 interval 轮迭代进行一次保存一次权重
   
   # 根据自己的 GPU 情况，修改 base_lr，修改的比例是 base_lr_default * (your_bs / default_bs)
   base_lr = _base_.base_lr / 8
   
   class_name = ('cat', )  # 根据 class_with_id.txt 类别信息，设置 class_name
   num_classes = len(class_name)
   metainfo = dict(
       classes=class_name,
       palette=[(220, 20, 60)]  # 画图时候的颜色，随便设置即可
   )
   
   train_cfg = dict(
       max_epochs=max_epochs,
       val_begin=20,  # 第几个 epoch 后验证，这里设置 20 是因为前 20 个 epoch 精度不高，测试意义不大，故跳过
       val_interval=save_epoch_intervals,  # 每 val_interval 轮迭代进行一次测试评估
       dynamic_intervals=[(max_epochs - _base_.num_last_epochs, 1)]
   )
   
   model = dict(
       bbox_head=dict(
           head_module=dict(num_classes=num_classes)),
       train_cfg=dict(
           initial_assigner=dict(num_classes=num_classes),
           assigner=dict(num_classes=num_classes))
   )
   
   train_dataloader = dict(
       batch_size=train_batch_size_per_gpu,
       num_workers=train_num_workers,
       dataset=dict(
           _delete_=True,
           type='RepeatDataset',
           # 数据量太少的话，可以使用 RepeatDataset ，在每个 epoch 内重复当前数据集 n 次，这里设置 5 是重复 5 次
           times=5,
           dataset=dict(
               type=_base_.dataset_type,
               data_root=data_root,
               metainfo=metainfo,
               ann_file='annotations/trainval.json',
               data_prefix=dict(img='images/'),
               filter_cfg=dict(filter_empty_gt=False, min_size=32),
               pipeline=_base_.train_pipeline)))
   
   val_dataloader = dict(
       dataset=dict(
           metainfo=metainfo,
           data_root=data_root,
           ann_file='annotations/trainval.json',
           data_prefix=dict(img='images/')))
   
   test_dataloader = val_dataloader
   
   val_evaluator = dict(ann_file=data_root + 'annotations/trainval.json')
   test_evaluator = val_evaluator
   
   optim_wrapper = dict(optimizer=dict(lr=base_lr))
   
   default_hooks = dict(
       # 设置间隔多少个 epoch 保存模型，以及保存模型最多几个，`save_best` 是另外保存最佳模型（推荐）
       checkpoint=dict(
           type='CheckpointHook',
           interval=save_epoch_intervals,
           max_keep_ckpts=5,
           save_best='auto'),
       param_scheduler=dict(max_epochs=max_epochs),
       # logger 输出的间隔
       logger=dict(type='LoggerHook', interval=10))
   
   custom_hooks = [
       dict(
           type='EMAHook',
           ema_type='ExpMomentumEMA',
           momentum=0.0001,
           update_buffers=True,
           strict_load=False,
           priority=49),
       dict(
           type='mmdet.PipelineSwitchHook',
           switch_epoch=max_epochs - _base_.num_last_epochs,
           switch_pipeline=_base_.train_pipeline_stage2)
   ]
   

2. 下载 YOLOv6-s 的预训练权重

   wget https://download.openmmlab.com/mmyolo/v0/yolov6/yolov6_s_syncbn_fast_8xb32-400e_coco/yolov6_s_syncbn_fast_8xb32-400e_coco_20221102_203035-932e1d91.pth -P work_dirs/

3. 训练

   python tools/train.py configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py

7. 推理

使用最佳的模型进行推理，下面命令中的最佳模型路径是 ./work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_96.pth，请用户自行修改为自己训练的最佳模型路径。

python demo/image_demo.py ./data/cat/images \
                          ./configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py \
                          ./work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_96.pth \
                          --out-dir ./data/cat/pred_images

小技巧

如果推理结果不理想，这里举例 2 种情况：

1. 模型欠拟合： 需要先判断是不是训练 epoch 不够导致的欠拟合，如果是训练不够，则修改 config 文件里面的 max_epochs 和 work_dir 参数，或者根据上面的命名方式新建一个 config 文件，重新进行训练。

2. 数据集需优化： 如果 epoch 加上去了还是不行，可以增加数据集数量，同时可以重新检查并优化数据集的标注，然后重新进行训练。

8.部署

MMYOLO 提供两种部署方式：

1. MMDeploy 框架进行部署

2. 使用 projects/easydeploy 进行部署

8.1 MMDeploy 框架进行部署

考虑到部署的机器环境千差万别，很多时候在本地机器可以，但是在生产环境则不一定，这里推荐使用 Docker，做到环境一次部署，终身使用，节省运维搭建环境和部署生产的时间。

 8.1.1 构建 Docker 镜像

git clone -b dev-1.x https://github.com/open-mmlab/mmdeploy.git
cd mmdeploy
docker build docker/GPU/ -t mmdeploy:gpu --build-arg USE_SRC_INSIDE=true

其中 USE_SRC_INSIDE=true 是拉取基础进行之后在内部切换国内源，构建速度会快一些。

执行脚本后，会进行构建，此刻需要等一段时间：

  8.1.2 创建 Docker 容器

export MMYOLO_PATH=/path/to/local/mmyolo # 先将您机器上 MMYOLO 的路径写入环境变量
docker run --gpus all --name mmyolo-deploy -v ${MMYOLO_PATH}:/root/workspace/mmyolo -it mmdeploy:gpu /bin/bash

 可以看到本地的 MMYOLO 环境已经挂载到容器里面了

8.1.3 转换 TensorRT 模型 

首先需要在 Docker 容器里面安装 MMYOLO 和 pycuda：

export MMYOLO_PATH=/root/workspace/mmyolo # 镜像中的路径，这里不需要修改
cd ${MMYOLO_PATH}
export MMYOLO_VERSION=$(python -c "import mmyolo.version as v; print(v.__version__)")  # 查看训练使用的 MMYOLO 版本号
echo "Using MMYOLO ${MMYOLO_VERSION}"
mim install --no-cache-dir mmyolo==${MMYOLO_VERSION}
pip install --no-cache-dir pycuda==2022.2

 进行模型转换

cd /root/workspace/mmdeploy
python ./tools/deploy.py \
    ${MMYOLO_PATH}/configs/deploy/detection_tensorrt-fp16_dynamic-192x192-960x960.py \
    ${MMYOLO_PATH}/configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py \
    ${MMYOLO_PATH}/work_dirs/yolov6_s_syncbn_fast_1xb32-100e_cat/best_coco/bbox_mAP_epoch_96.pth \
    ${MMYOLO_PATH}/data/cat/images/mmexport1633684751291.jpg \
    --test-img ${MMYOLO_PATH}/data/cat/images/mmexport1633684751291.jpg \
    --work-dir ./work_dir/yolov6_s_syncbn_fast_1xb32-100e_cat_deploy_dynamic_fp16 \
    --device cuda:0 \
    --log-level INFO \
    --show \
    --dump-info

等待一段时间，出现了 All process success. 即为成功：

 查看导出的路径，可以看到如下图所示的文件结构：

$WORK_DIR
  ├── deploy.json
  ├── detail.json
  ├── end2end.engine
  ├── end2end.onnx
  └── pipeline.json

8.1.4 部署模型执行推理

需要将 ${MMYOLO_PATH}/configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py 里面的 data_root 修改为 Docker 容器里面的路径：

data_root = '/root/workspace/mmyolo/data/cat/'  # Docker 容器里面数据集目录的绝对路径

执行速度和精度测试：

python tools/test.py \
    ${MMYOLO_PATH}/configs/deploy/detection_tensorrt-fp16_dynamic-192x192-960x960.py \
    ${MMYOLO_PATH}/configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py \
    --model ./work_dir/yolov6_s_syncbn_fast_1xb32-100e_cat_deploy_dynamic_fp16/end2end.engine \
    --speed-test \
    --device cuda

速度测试如下，可见平均推理速度是 24.10 ms，对比 PyTorch 推理有速度提升，同时显存也下降了很多：

Epoch(test) [ 10/116]    eta: 0:00:20  time: 0.1919  data_time: 0.1330  memory: 12
Epoch(test) [ 20/116]    eta: 0:00:15  time: 0.1220  data_time: 0.0939  memory: 12
Epoch(test) [ 30/116]    eta: 0:00:12  time: 0.1168  data_time: 0.0850  memory: 12
Epoch(test) [ 40/116]    eta: 0:00:10  time: 0.1241  data_time: 0.0940  memory: 12
Epoch(test) [ 50/116]    eta: 0:00:08  time: 0.0974  data_time: 0.0696  memory: 12
Epoch(test) [ 60/116]    eta: 0:00:06  time: 0.0865  data_time: 0.0547  memory: 16
Epoch(test) [ 70/116]    eta: 0:00:05  time: 0.1521  data_time: 0.1226  memory: 16
Epoch(test) [ 80/116]    eta: 0:00:04  time: 0.1364  data_time: 0.1056  memory: 12
Epoch(test) [ 90/116]    eta: 0:00:03  time: 0.0923  data_time: 0.0627  memory: 12
Epoch(test) [100/116]    eta: 0:00:01  time: 0.0844  data_time: 0.0583  memory: 12
[tensorrt]-110 times per count: 24.10 ms, 41.50 FPS
Epoch(test) [110/116]    eta: 0:00:00  time: 0.1085  data_time: 0.0832  memory: 12

精度测试如下。此配置采用 FP16 格式推理，会有一定程度掉点，但是推理速度更快、显存占比更小：

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.954
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 1.000
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.975
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.954
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.860
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.965
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.965
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.965

INFO - bbox_mAP_copypaste: 0.954 1.000 0.975 -1.000 -1.000 0.954
INFO - Epoch(test) [116/116]  coco/bbox_mAP: 0.9540  coco/bbox_mAP_50: 1.0000  coco/bbox_mAP_75: 0.9750  coco/bbox_mAP_s: -1.0000  coco/bbox_mAP_m: -1.0000  coco/bbox_mAP_l: 0.9540

部署模型图片推理演示：

cd ${MMYOLO_PATH}/demo
python deploy_demo.py \
    ${MMYOLO_PATH}/data/cat/images/mmexport1633684900217.jpg \
    ${MMYOLO_PATH}/configs/custom_dataset/yolov6_s_syncbn_fast_1xb32-100e_cat.py \
    /root/workspace/mmdeploy/work_dir/yolov6_s_syncbn_fast_1xb32-100e_cat_deploy_dynamic_fp16/end2end.engine \
    --deploy-cfg ${MMYOLO_PATH}/configs/deploy/detection_tensorrt-fp16_dynamic-192x192-960x960.py \
    --out-dir ${MMYOLO_PATH}/work_dirs/deploy_predict_out \
    --device cuda:0 \
    --score-thr 0.5

执行之后，可以看到在 --out-dir 下面的推理图片结果

8.1.5 保存和加载 Docker 容器 

因为如果每次都进行 docker 镜像的构建，特别费时间，此时您可以考虑使用 docker 自带的打包 api 进行打包和加载。

# 保存，得到的 tar 包可以放到移动硬盘
docker save mmyolo-deploy > mmyolo-deploy.tar

# 加载镜像到系统
docker load < /path/to/mmyolo-deploy.tar

8.1.6 进行部署

使用projects/easydeploy部署 下方转到官方文档

mmyolo/projects/easydeploy/README_zh-CN.md at dev · open-mmlab/mmyolo · GitHub

难点：

1. 针对训练与预测的图像分辨率，进行参数设置

在要训练的模型代码中例如：yolov5_s-v61_fast_1xb12-40e_608x352_xxx.py 加上以下代码

# 根据实际图片分辨率进行更改
img_scale = (640, 480)
# 在这部分之前
max_epochs = 200
train_batch_size_per_gpu = 14
train_num_workers = 2

2. 针对很小目标如何检测或者调参优化(仍在研究)