本文介绍干货：SOLOv2训练自己的数据集保姆级教程附代码！

干货：SOLOv2训练自己的数据集保姆级教程附代码！

全文干货较多，建议收藏再看，顺便点个赞！

本篇文章主要介绍如何使用SOLOv2训练自己的数据集．在这之前有人已经能够把SOLOv2转到TensorRT进行工业化部署了．那么如何训练自己的数据集是一个前置问题，鉴于训练数据集很多人都有需求，本篇文章就教大家如何一步步的实现，同时我们也会在文章最下方放上代码链接．如果你觉得有用，欢迎分享，转发，并关注我们的公众号！

请注意，本篇文章不是基于官方的mmdetection版本！请下载我们fork(文末有链接)版本进行操作．

总结一下我们的版本包含的新增地方：

自定义数据集的训练
SOLOv2-Lite版本
可以更换backbone
支持ONNX export
TensorRT加速（内部）

SOLOv2的牛逼之处

最近看PaddlePaddle detection 发布了2.0版本，我看了下里面的SOlOv2的效果，确实很不错了啊

但是奈何我不怎么用Paddle. 感觉坑应该也不浅，所以就想用pytorch做这么一套．

实际搞下来我感觉我们做的也还不错：

我们不仅可以出框，还可以出mask, 这五颜六色花花绿绿的着实漂亮．需要注意的是上面那一张图是纯C++推理的结果，而且可以使用TensorRT加速．也就是说，我们已经把训练和部署这一套打通了．

框架概览

我们使用的不是官方solov2的mmdetection版本，官方那个太老旧了，还是基于mmdetection1.0，虽然我们之前依据2.0的升级方法升级了一个版本，但是感觉也不是特别好用．这回切换到了detectron2.

后面的教程都会基于detectron2来讲．

如何准备自己的数据集

我们使用的数据集格式是coco. 因为这个数据格式用的比较多，在dectron2里面支持也比较好．假如你的数据已经是coco格式了，那么很好，你只需要进行可视化一下就可以了．

这里给大家安利一个工具：alfred-py. 这是一位大佬开源的深度学习工具集合，安装方式很简单：

pip install alfred-py

然后我们就可以使用alfred 命令可视化你的coco数据．

怎么做呢？

alfred data cocoview -j annotations.json -i images/

传入你的json和图片的根目录，可视化的结果就出来了．不信的你可以直接对coco数据可视化试一下．

这个工具在验证你的数据格式的时候还是很有用的！同时也很方便的让我们对数据进行检查

这种好工具难道不应该分享给身边的人用？简直是一颗关键时候可以用的银弹．

初次之外这个库还有很多其他的功能，大家可以-h去慢慢的探索．这个工具还可以实现不同数据集之间的转换！比如coco转到yolo, voc转到yolo等，全部是一行命令完成！不需要编写任何代码！

如果你的数据没问题了，我们就进行下一步．

我们的实例数据集

这里给大家暂时一下我们自己的数据集．我们虽然在coco上已经成功训练了．最好还是在一个自己的数据上训练一下，我们采用的数据集是一个简单的微控制器分类，主要分为:

CLASS_NAMES = ["Arduino", 'ESP8266',  'Heltec', 'Raspberry']

大概长这样：

这个数据集是开源的，大家可以去kaggle上或者google搜索一下，微控制器实力分割．

注册自己的数据集

在d2里面，最好的训练自己的数据集的方式就是把自己的dataset注册进入d2. 这样你就可以在任何地方调用你的类别．我们的方法通常是在 tools 文件夹下拷贝一份　train_net.py.

同时做这个更改：

CLASS_NAMES = ["Arduino", 'ESP8266',  'Heltec', 'Raspberry']
DATASET_ROOT = './data/microship'
ANN_ROOT = os.path.join(DATASET_ROOT, 'annotations')
TRAIN_PATH = os.path.join(DATASET_ROOT, 'train')
VAL_PATH = os.path.join(DATASET_ROOT, 'test')
TRAIN_JSON = os.path.join(ANN_ROOT, 'train.json')
VAL_JSON = os.path.join(ANN_ROOT, 'test.json')
PREDEFINED_SPLITS_DATASET = {
    "microship_train": (TRAIN_PATH, TRAIN_JSON),
    "microship_val": (VAL_PATH, VAL_JSON),
}


def plain_register_dataset():
    # 训练集
    DatasetCatalog.register(
        "microship_train", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH))
    MetadataCatalog.get("microship_train").set(thing_classes=CLASS_NAMES,
                                            evaluator_type='coco',
                                            json_file=TRAIN_JSON,
                                            image_root=TRAIN_PATH)
    DatasetCatalog.register(
        "microship_val", lambda: load_coco_json(VAL_JSON, VAL_PATH))
    MetadataCatalog.get("microship_val").set(thing_classes=CLASS_NAMES,
                                          evaluator_type='coco',
                                          json_file=VAL_JSON,
                                          image_root=VAL_PATH)


plain_register_dataset()

简单来说做这些指定：

你的类别名称；
你的标注，数据的存放路径，分train 和　val.

然后注册你的数据集名称．

最后在　configs下面，在这个配置内：

MODEL:
  META_ARCHITECTURE: "SOLOv2"
  MASK_ON: True
  BACKBONE:
    NAME: "build_resnet_fpn_backbone"
  RESNETS:
    OUT_FEATURES: ["res2", "res3", "res4", "res5"]
  FPN:
    IN_FEATURES: ["res2", "res3", "res4", "res5"]
DATASETS:
  TRAIN: ("microship_train",)
  TEST: ("microship_val",)
SOLVER:
  IMS_PER_BATCH: 4
  BASE_LR: 0.01
  WARMUP_FACTOR: 0.01
  WARMUP_ITERS: 1000
  STEPS: (60000, 80000)
  MAX_ITER: 90000
INPUT:
  MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)
  MASK_FORMAT: "bitmask"
VERSION: 2

写上你的数据集名字就可以了．

在这里需要注意一个地方：

如果你是单卡训练，这里的　BASE_LR 要除以８；
因为默认是８卡

我个人建议如果是非常大的数据集，还是使用８卡徐莲会收敛快一点．

如何发起训练

发起训练就一行命令了：

python3 tools/train_microship.py --config-file configs/SOLOv2/microship_dataset/R50_lite_3x.yaml

把这个config传入就可以了．

对了，别忘了确保你注册的数据集被正确的软链接．

好了，本次干货分享就到这儿，需要代码的可以在这里获取：

manaai.cn/aicodes_det…