Yolox模型架构
模型结构由backbonex、neck、head组成
backbone:CSPDarknet53
neck:PAFPN
head:Decoupled Head、Anchor Free、SimOTA
1. Yolox 环境搭建
根目录yolox
1.1 创建并激活环境
使用的环境为conda的yolox, python 3.8
激活环境
conda activate yolox
1.2 安装代码库的文件
切换到根目录下,终端上
pip install -r requirements.txt
python3 setup.py develop
1.3 下载安装apex
git@github.com:NVIDIA/apex.git
切换到apex目录下,安装
python3 setup.py install
1.4 下载pycocotools
pip3 install cython
pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
2. 创建自己的数据集
VOC数据集
目录
VOCdevkit # 放在./datasets下
___VOC2007
_____Annotations(存放所有图片对应的xml标签)
_____ImageSets (划分训练集、验证集、测试集的txt文件)
_____JPEGImages (存放所有图片,格式为png/jpg)
训练集和测试集生成的代码如下
# coding = utf-8
# -*- coding:utf-8 -*-
import os
import random
xml_file_path = '/home/zw100/Deep Learning/yolox/datasets/VOCdevkit/VOC2007/Annotations'
txt_save_path = '/home/zw100/Deep Learning/yolox/datasets/VOCdevkit/VOC2007/ImageSets/Main'
def generate_trainval_test(xml_file_path, txt_save_path):
test_percent = 0.1
xml_files = os.listdir(xml_file_path)
num_of_xml_files = len(xml_files)
test_num = int(num_of_xml_files * test_percent)
list = range(num_of_xml_files)
test = random.sample(list, test_num)
f_test = open(txt_save_path+'/test.txt', 'w')
f_train = open(txt_save_path+'/trainval.txt', 'w')
for i in list:
name = xml_files[i][:-4] + '\n' # q去掉.xml,即[:-4]
if i in test:
f_test.write(name)
else:
f_train.write(name)
f_train.close()
f_test.close()
if __name__ == '__main__':
generate_trainval_test(xml_file_path, txt_save_path)
3. 修改代码
3.1 exps目录
./exps/example/yolox_voc/yolox_voc_s.py
self.num_classes --->你自己的目标检测的类别数,本文为检测猫狗,所以self.num_classes = 2.
3.2 yolox目录
yolox/data/datasets/voc_classes.py
VOC_CLASSES--->改成自己要检测的目标,本文为猫狗检测,所以
VOC_CLASSES={"cat", "dog"}
voc.py
self._imgpath根据自己的图片是png还是jpg格式更改,本文是png,所以改为
self._imgpath = os.path.join("%s", "JPEGImages", "s%.png")
4. 训练
需要在yolox的官网下载yolox_s.pth预训练权重文件,Yolox github官网
python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0-b 4--fp16 -c yolox_s.pth
-c后面接yolox_s.pth的路径,本文直接放在根目录下
5. 测试
直接测试图片集,即test_txt对应的图片
利用tools下的eval.py进行测试
修改eval.py,将-c对应的pth改成训练集训练出来的best_ckpt.pth路径,如
如:
python tools/eval.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0 -b 4 --fp16 -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth
测试单个图片
利用tools下的demo.py进行测试
python tools/demo.py image -f exps/example/yolox_voc/yolox_voc_s.py -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth --path assets/class01.jpg --conf 0.25--nms 0.45--tsize 640--save_result --device [cpu/gpu]
6. 部署
yolox是先通过tools/trt.py脚本把ONNX模型解析然后序列化到model_trt.engin文件中,最后再通过c++代码从该文件中加载模型去进行推理
可以通过python解析onnx
也可以通过c++解析onnx
6.1 python版tensorRT推理
运行tools/trt.py,会在实验结果output(如YOLOX_outputs/yolox_s_add_kaiyang-736-1312-125-new3)下生成model_trt.engin
并且复制model_trt.engine到demo/cpp(c++)下
python tools/trt.py -f /path/to/your/yolox/exps/yolox_s.py -c your_ckpt.pth #example
python tools/trt.py -f exps/example/yolox_voc/yolox_voc_s_service_area_add_kaiyang.py -c YOLOX_outputs/yolox_s_addkaiyang-736-1312-125-new3/best_ckpt.pth
然后运行demo推理图片
python tools/demo.py image -f exps/default/yolox_s.py --trt --save_result# example
python tools/demo.py image -f exps/example/yolox_voc/yolox_voc_s_service_area_add_kaiyang.py --path datasets/VOCdevkit/VOC2007/TestImage --trt --save_result
infer_time = 0.003左右
6.2 c++版tensorRT推理
切换到demo/TensorRT/cpp目录下
编写CMakeLists.txt,包括cuda路径,TensorRT路径
新建bulid,并cmake
cd build
cmake ..
make
若在此过程编译失败例如找不到CUDA,卸载重装cuda其他不变动(cuda)
运行demo
./yolox ../model_trt.engine -i ../../../../assets/dog.jpg
./yolox ../model_trt.engine -i ../../../../datasets/VOCdevkit/VOC2007/TestImage
./yolox ../model_trt.engine -i ../../../../datasets/VOCdevkit/VOC2007/JPEGImages/1.jpg<