yolact_edge训练自己的yolact_edge模型(并部署在Jetson Xavier上)

Posted Rainbowman 0

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了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,以上就是整个过程。这篇博客只是记录我的训练和部署过程,大致过程不会出问题,具体细节有问题需要自行百度解决。

以上是关于yolact_edge训练自己的yolact_edge模型(并部署在Jetson Xavier上)的主要内容,如果未能解决你的问题,请参考以下文章

paddleocr训练自己的数据最简单方式软件一键训练

深度学习训练营之yolov5训练自己的数据集

怎样用自己的数据集对caffe训练好的model进行fineture

图像识别之Yolov5训练自己的模型

图像识别之Yolov5训练自己的模型

Caffe 训练和测试自己的图片