Nanodet模型的训练以及移植到ROS

Nanodet模型的训练以及移植到ROS

Nanodet模型，相对yolo来说更加轻量化
为了配合后面的ros-nanodet需要选择v0.3.0版本

模型训练：Nanodet
Ros调用Nanodet：ros-nanodet

数据集准备

使用的是COCO格式数据集，需要使用labelme工具做标注

1
2
pip install labelme
labelme
  1. 打开图片目录
  2. 把每一个图片目标框选出来并保存
  3. 保存的格式为.json，每一张图片对应一个json

标注完了之后需要将所有的.json格式的文件转为COCO格式

1
2
pip install labelme2coco
labelme2coco /path/to/labelme_json_dir --output /path/to/output.json

output.json文件就是COCO格式的标注
大致格式如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
{
  "images": [
    {
      "id": 1,
      "file_name": "000001.jpg",
      "width": 1280,
      "height": 720
    },
    {
      "id": 2,
      "file_name": "000002.jpg",
      "width": 1280,
      "height": 720
    }
  ],
  "annotations": [
    {
      "id": 1,
      "image_id": 1,
      "category_id": 1,
      "bbox": [100, 200, 300, 400],
      "area": 120000,
      "iscrowd": 0
    },
    {
      "id": 2,
      "image_id": 2,
      "category_id": 2,
      "bbox": [50, 60, 200, 300],
      "area": 60000,
      "iscrowd": 0
    }
  ],
  "categories": [
    {"id": 0, "name": "cat"},
    {"id": 1, "name": "dog"}
  ]
}

文件目录：
因为我的图片文件夹对train和val做区分，所以在COCO的标注数据中，要在文件名前面加上val\\和train\\

1
2
3
4
5
6
7
8
9
10
images
├── train
│   ├── train_pic1.jpg
│   └── train_pic2.jpg
└── val
    ├── val_pic1.jpg
    └── val_pic2.jpg
annotations
├── train.json
└── val.json

训练环境

环境:

  • python：3.8.20
  • Cuda：11.8
  • pytorch：2.4.1+cu118

配置conda环境

1
pip install -r requirements.txt -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple

然后这里安装完了之后需要删除pytorch及其插件，安装cuda版本的pytorch

1
2
pip uninstall torch torchvision torchaudio
pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu118

安装NanoDet

1
python setup.py develop

训练配置文件

修改config/nanodet-m.yml

修改内容

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 根据类别数设置
num_classes: 2

# 根据图片路径设置
train:
    img_path: /data/images
    ann_path: /data/annotations/train.json
val:
	img_path: /data/images
    ann_path: /data/annotations/val.json

# 根据设备情况设置
  workers_per_gpu: 1
  batchsize_per_gpu: 8
# 训练的轮数  
  total_epochs: 300
  
# 根据数据标注的名称设置
class_names: ['Class1', 'Class2']

训练

1
python tools/train.py config/nanodet-m.yml

训练的中的报错

问题一

1
Traceback (most recent call last): File "tools/train.py", line 24, in <module> from nanodet.data.collate import collate_function File "\nanodet-0.3.0\nanodet\data\collate.py", line 5, in <module> from torch._six import container_abcs, string_classes, int_classes ModuleNotFoundError: No module named 'torch._six'

这个错误是因为 PyTorch 版本过高导致的，torch._six 模块在较新的 PyTorch 版本（如 1.10 及以上）中已被移除，而你使用的代码（nanodet）可能是基于旧版本 PyTorch 编写的。

解决方案：
修改报错文件为新版的代码即可

  1. 找到nanodet/data/collate.py 的第 5 行
1
from torch._six import container_abcs, string_classes, int_classes
  1. 替换为下面的内容
1
2
3
import collections.abc as container_abcs import six 
string_classes = (str, bytes)
int_classes = (int,)
  1. pip下载six
1
pip install six

问题二

1
[root][10-18 12:55:38]INFO:Creating model... model size is 1.0x init weights... => loading pretrained model [https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth](https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth) Finish initialize Lite GFL Head. Traceback (most recent call last): File "tools/train.py", line 96, in <module> main(args) File "tools/train.py", line 88, in main callbacks=[ProgressBar(refresh_rate=0)] # disable tqdm bar TypeError: __init__() got an unexpected keyword argument 'refresh_rate'

这个错误是因为使用的 PyTorch Lightning 版本与代码中 ProgressBar 的参数不兼容导致的。在较新的 PyTorch Lightning 版本中，ProgressBar 类的初始化参数 refresh_rate 已被移除或更改（例如，在 PyTorch Lightning 2.0+ 中，进度条相关的参数调整较大）

解决方案：

  1. 找到tools/train中的ProgressBar类
1
callbacks=[ProgressBar(refresh_rate=0)]
  1. 修改为
1
callbacks=[TQDMProgressBar(refresh_rate=0)]
  1. 引入TQDMProgressBar
1
from lightning.pytorch.callbacks import TQDMProgressBar

问题三

1
Traceback (most recent call last): File "tools/train.py", line 96, in <module> main(args) File "tools/train.py", line 80, in main trainer = pl.Trainer(default_root_dir=cfg.save_dir, File "\nanodet-0.3.0\.conda\lib\site-packages\pytorch_lightning\utilities\argparse.py", line 348, in insert_env_defaults return fn(self, **kwargs) File "\nanodet-0.3.0\.conda\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 420, in __init__ self._accelerator_connector = AcceleratorConnector( File "nanodet-0.3.0\.conda\lib\site-packages\pytorch_lightning\trainer\connectors\accelerator_connector.py", line 183, in __init__ self._check_config_and_set_final_flags( File "\nanodet-0.3.0\.conda\lib\site-packages\pytorch_lightning\trainer\connectors\accelerator_connector.py", line 294, in _check_config_and_set_final_flags raise ValueError( ValueError: You selected an invalid accelerator name: `accelerator='ddp'`. Available names are: cpu, cuda, hpu, ipu, mps, tpu.

这个错误是因为 PyTorch Lightning 版本更新后，加速器（accelerator）的配置方式发生了变化。在旧版本中，accelerator='ddp' 是有效的分布式训练配置，但在新版本（如 PyTorch Lightning 2.0+）中，分布式训练策略（如 DDP）需要通过 strategy 参数单独设置，而 accelerator 参数仅用于指定硬件类型（如 CPU、GPU 等）

解决方案：

  1. 找到tools/train.py中的accelerator大概在84行
  2. 将ddp改为cuda

修改完了之后就不会再报错了

模型转换

训练完的模型是.ckpt需要转换为.onnx

pip下载onnx

1
pip install onnx

转换命令

1
python  .\tools\export.py --cfg_path .\config\nanodet-m.yml --model_path .\workspace\nanodet_m\model_best\model_best.ckpt --out_path .\workspace\nanodet_m\model_best\model_best.onnx

执行命令后会在.\workspace\nanodet_m\model_best\中生成model_best.onnx 模型

这是模型就已经处理完成

模型部署

1
2
3
mkdir -p nanodet_ros/src  
cd nanodet_ros/src  
git clone https://github.com/stunback/ros-nanodet.git

文件目录:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
├── CMakeLists.txt
├── package.xml
├── scripts
│   └── ros_nanodet.py
├── setup.py
└── src
    └── nanodet_ros
        ├── __init__.py
        ├── model
        │   ├── coco.names
        │   └── nanodet.onnx
        ├── nanodet.py
        └── __pycache__
            └── nanodet.cpython-38.pyc

这里需要把我们自己的onnx模型丢到model中，并且写一个新的.names文件

class.names

1
2
Com
NonCom

修改文件适配自己的模型

修改scripts/ros_nanodet.py

1
2
3
4
5
6
7
8
9
10
11
12
13
# 改成和训练的config一样的尺寸
# 18、19行
IMAGE_WIDTH = 320
IMAGE_HEIGHT = 320

# 改成自己的文件名
# 48、49行
model_path = os.path.join(model_dir, 'model/nanodet-3.0.onnx')
clsname_path = os.path.join(model_dir, 'model/class.names')

# 改成自己设置的摄像头的话题
# 59行
image_topic_1 = "/camera/image_raw"

测试

1
2
3
4
cd ~/nanodet_ros 
catkin_make
source devel/setup.sh
rosrun nanodet_ros ros_nanodet.py

测试成功！！！

ROS
#ROS #NanoDet
Nanodet模型的训练以及移植到ROS
http://www.ming-ice-tea.top/2025/10/19/Nanodet模型的训练以及移植到ROS中/
作者
Ming
发布于
2025年10月19日
许可协议