nanodet-plus训练自己数据集

序言

前两天nanodet-plus隆重发布，又赚了一波热度，趁着年底有空，赶紧学习一波，避免被卷死；因为之前有过nanodet的训练实践经历，但是有好长一段时间没用了，代码都生疏了，还好作者将新版本合并到老仓库中，代码结构基本上没变，旧的配置文件修改依旧适用，重新上手起来也比较容易，这次打算拿之前的筷子数据集来试试效果，本文记录训练过程。

NanoDet-Plus与上一代NanoDet相比，在仅增加1毫秒多的延时的情况下，精度提升了30%，与YOLOv5-n, YOLOX-Nano等其他轻量级模型相比，在精度和速度上也都高了不少！同时NanoDet-Plus改进了代码和架构，提出了一种非常简单的训练辅助模块，使模型变得更易训练！同时新版本也更易部署，同时提供ncnn、OpenVINO、MNN以及安卓APP的Demo！

先看下nanodet-plus与上一代nanodet以及其他轻量型的检测算法对比：

具体改进请参考作者知乎原文介绍：Nanodet-Plus 知乎介绍

一、环境配置

简单介绍一下我的本机环境：

• ubuntu 18.04
• RTX 3070
• cuda 11.2
• python 3.7
• pytorch 1.8

1.1 nanodet安装

需要安装nanodet及相关依赖，打开终端界面，在终端中输入命令如下：

git clone https://github.com/RangiLyu/nanodet.git
cd nanodet
pip install -r requirements.txt
python setup.py develop

安装成功后进行测试，在nanodet文件夹下创建一个images文件夹，将测试图片放入该文件夹中，然后下载coco数据集训练的权重（权重放在谷歌云盘上，下载的话需要科学上网），运行测试demo：

python demo/demo.py image --config CONFIG_PATH --model MODEL_PATH --path IMAGE_PATH

• CONFIG_PATH 是nanodet-plus的yml配置文件
• MODEL_PATH 是刚才下载的权重路径
• IMAGE_PATH 是测试图片路径

权重下载连接地址，可以将weight和checkpoint都下载下来，推理的时候只使用weight，训练的时候会用到checkpoint，我这里只下了416尺寸的：

以下是我的测试例子，下载的权重放在新建的weights文件夹中，使用nanodet-plus-m的416尺寸推理：

 python demo/demo.py image --config config/nanodet-plus-m_416.yml --model weights/nanodet-plus-m_416.pth --path images/

测试结果如下，误检率还是挺高的，不过置信度都比较低，可以通过阈值过滤掉：

或者使用nanodet-plus-m-1.5x_416版本，效果会好很多：

好了，先不纠结精度的问题，测试成功的话说明环境已经安装成功了，如果运行出错的话需要自己排查一下，按照步骤来一般不会有什么问题。

二、数据准备

为了以后方便测试某个新出的目标检测算法，介于coco和voc数据集都太大了，训练一次要耗费很长的时间，所以我自己构建了一个筷子数据集，数量只有两百多张，图片少了点，但是非常方便训练配置，数据集小，训练速度快，很容易出结果。具有很好的模型测试意义，具体介绍请看我之前的构建文章构建了一个用于目标检测点数的数据集

然后根据配置文件数据集读入要求，将voc格式的数据集分为train和val两部分，文件夹划分如下，因为数据集数据量少，不想写脚本划分的话，自己手动复制粘贴划分一下也没问题，记录下路径：

三、修改配置文件

将config/nanodet_custom_xml_dataset.yml复制一份，另取名为config/nanodet_plus-m-416_test.yml，按需修改以下几处：

训练保存的文件夹：

训练目标类别数，根据自己数据集来，我的只有一类：

数据集标签类别：

训练尺寸和数据集路径：

依次为训练批次大小（根据自己的显卡来）、预训练权重路径（好像不加收敛速度也很快）、训练总轮次：

修改结束后运行训练命令：

python tools/train.py config/config/nanodet_plus-m-416_test.yml

出现如下训练界面说明模型已经开始训练，如果遇到问题的话，检查一下自己是不是配置文件没修改成功，具体问题具体排查：

最终训练精度结果如下：

还不错的样子，AP0.5:0.95达到了 0.749，AP0.5 也达到了0.989，具体看可视化效果如何。

运行测试demo，可视化结果如下，效果好像稍微差一丢丢，不过问题不大，可以通过调整后处理阈值过滤，或者训练过程调整一些超参数来更好的训练，因为只是简单跑通看看效果，结果实际上看起来还是很不错的，上1.5x的模型效果应该更好：