【yolact_edge】训练自己的yolact_edge模型（并部署在Jetson Xavier上）

0. 前言

前段时间在Jetson Xavier上试了试官方的demo，发现效果还不错，帧率能跑到13fps左右，昨天开始用自己的数据集训练yolact_edge模型，记录下大致过程。

1. 整个流程

整个流程比较简单，大致如下：

（1）在Jetson Xavier上部署基本环境
（2）按coco格式准备自己的数据集
（3）修改yolact_edge/data/config.py文件的配置
（4）训练模型，检测模型

2. 具体过程

2.1 部署基本环境

基本环境的部署我上一篇写过，参考这里。

2.2 按COCO格式准备自己的数据集

这里的大致流程为：
（1）安装labelme工具
（2）使用labelme工具打标签
（3）将数据转换成coco格式
（4）将数据按文件夹分为train、val、test
具体过程我在之前的博客也写有，参考这篇博客的2.2和2.3部分（我这篇博客写的很清楚了）。

2.3 修改yolact_edge/data/config.py文件的配置

我们在使用train.py训练模型时，相关配置都放在yolact_edge/data/config.py文件中，我们要根据自己的数据集来修改配置文件，具体需要修改以下几处：

2.3.1 修改一：增加自己的数据集定义

我们需要自己定义我们数据集的种类和标签，比如我的标签只有一类“truck”，就修改如下：

2.3.2 修改二：在DATASET部分增加自己的数据集

其中：
（1）name：你的数据集名字
（2）class_name：你的类的名称，就是你在2.3.1部分起的类名
（3）label_map：你的标签名称，就是你在2.3.1部分起的标签名称
（4）train_info：你在2.2部分训练集数据通过labelme2coco.py文件生成的.json文件所在位置
（5）train_images：你的训练集数据所在位置
（6）valid_info：你在2.2部分验证集数据通过labelme2coco.py文件生成的.json文件所在位置
（7）valid_images：你的验证集数据所在位置

2.3.3 修改三：增加自己数据集的训练配置

（1）name：自己定义的名字
（2）dataset：你在2.3.2中的dataset
（3）num_classes：在类别数的基础上要加1
（4）backbone：主干网络，写我写的这个就行
（5）max_iter：最大迭代次数，我写的50是因为我想先做一个初步的尝试。

2.3.4 修改四：在代码的最后制定cfg为你2.3.3配置的cfg名称

2.4 训练模型，检测模型

首先需要下载resnet50-19c8e357.pth文件作为基模型，下载链接我放在下面：

链接：https://pan.baidu.com/s/14zdlkneSjAClFmS-DQ4lQg
提取码：aaaa

下载好文件后新建一个weights的文件夹，把resnet50-19c8e357.pth放在这个文件夹下，然后运行命令：

python train.py --config=你的config名称（即你在2.3.4中的cfg名称）

例如我的命令是：

python train.py --config=yolact_edge_truck_config

然后就开始了训练过程，训练可能会提前结束，训练好的模型会存在weights文件夹中。我们将训练好的.pth文件拷到Jetson Xavier上，然后进行测试：

python3 eval.py --trained_model=weights/你训练的模型 --score_threshold=0.3 --top_k=100 --video=0 --ues_tensorrt_safe_mode

例如我的命令是：

python3 eval.py --trained_model=weights/yolact_edge_truck_25_50.pth --score_threshold=0.3 --top_k=100 --video=0 --use_tensorrt_safe_mode

这里注意，一定要加–use_tensorrt_safe_mode，否则会报错，类似于：IndexError:The shape of mask [32] at index 0 does not match the shape of the indexed tensor [1, 32] at index 0

测试的结果还可以，帧率能达到13fps左右。

3. 总结

OK，以上就是整个过程。这篇博客只是记录我的训练和部署过程，大致过程不会出问题，具体细节有问题需要自行百度解决。