记一次训练 YOLOX 模型

数据标注

使用 labelme 进行数据标注。

打开目录: 打开图片所在目录。
更改输出路径: 改变标注的 JSON 文件输出目录, 默认和图片目录一致。
自动保存：勾选后将会在标注后自动保存, 否则每切换一张图片会弹出提示框确认

加载图片后，点击创建多边形即可开始标注。点击顶栏的编辑或者在图片上右键，可以改变标注类型。

这里我使用的是创建矩形。

处理数据集

编写 labels.txt 文件，前两行固定内容，后面一行一个分类标签

__ignore__
_background_
英文分类名称

提前把图片分成两组，一组 train，一组 val，图片和JSON一起。比例可以是 0.85:0.15

使用 labelme2coco.py 分别转换为 COCO 格式。请确输入目录中同时包含图像和JSON文件。

python labelme2coco.py --labels labels.txt "输入目录" "输出目录"

然后可以看到目录结构

输出目录
├─JPEGImages
│  └─图片...
├─Visualization
│  └─可视化结果...
└─annotations.json

训练

安装YOLOX

克隆 https://github.com/Megvii-BaseDetection/YOLOX 仓库。

确保已经安装好了 PyTorch。

执行 pip3 install -v -e . 安装 yolox。注意路径中不要有中文、空格!!

如果安装时报错 [WinError 206] 文件名或扩展名太长。, 打开注册表编辑器，定位到HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Control\FileSystem，修改LongPathsEnabled为1即可。

我安装 onnx-simplifier 依赖时，报错装不上。于是我把 onnx-simplifier==0.4.10 修改为了 onnx-simplifier>=0.4.10 就装上了。

放置数据集

YOLOX 安装成功后，进入 datasets 目录，如下放置数据集

datasets
└─COCO
    ├─annotations
    │      instances_train2017.json
    │      instances_val2017.json
    ├─train2017
    └─val2017

train2017 和 val2017 目录存放图片。

准备预训练模型

下载 yolox_s.pth 放到根目录。

修改exp文件

到 exps/default 目录下，找到预训练模型对应的文件，在其中设置分类数量。

一般是你的分类数+1。具体可以到 instances_xxxx2017.json 中查看 categories 的成员数。

以 yolox_s.py 为例:

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 0.33
        self.width = 0.50
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        self.num_classes = 1

修改 yolox/data/dataset/coco_classes.py 文件，确保其中定义的类别和instances_xxxx2017.json里面categories定义的一致。

COCO_CLASSES = (
    "_background_",
    "你的分类名称",
)

开始训练

# -d 表示使用的设备
# -b 表示批处理规模
# 如果你的设备性能良好, 可以适当增大这两个值; 如果内存不足报错 OOM 则需要减小
# 官方的示例甚至是 -d 8 -b 64，而我的设备 -d 1 -b 8 就是极限了 T_T
python tools/train.py -f exps/default/yolox_s.py -d 1 -b 8 --fp16 -o -c yolox_s.pth

如果训练到一半中断了，可以继续训练

# 注意这里的模型要使用最后保存的模型
python tools/train.py -f exps/default/yolox_s.py -d 1 -b 8 --fp16 -o -c "YOLOX_outputs/yolox_s/latest_ckpt.pth" --resume

中途报错处理

ModuleNotFoundError: No module named 'yolox.layers.fast_cocoeval'

搜索 fast_cocoeval 修改为 fast_coco_eval_api 即可

yolox/layers/jit_ops.py

# def __init__(self, name="fast_cocoeval"):
    def __init__(self, name="fast_coco_eval_api"):
        super().__init__(name=name)

AttributeError: module 'yolox.layers.fast_coco_eval_api' has no attribute 'InstanceAnnotation'

搜索 from yolox.layers import COCOeval_opt as COCOeval 修改为 from pycocotools.cocoeval import COCOeval

yolox/evaluators/coco_evaluator.py

try:
    # from yolox.layers import COCOeval_opt as COCOeval
    from pycocotools.cocoeval import COCOeval
except ImportError:
    from pycocotools.cocoeval import COCOeval

推理

# -n 模型配置
# 本文中都以 yolox-s 为例
python tools/demo.py image -n yolox-s -c "YOLOX_outputs/yolox_s/best_ckpt.pth" --path "assets/1.png" --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

导出onnx模型

导出

先安装 onnxruntime，确保版本支持你的环境。

# -n 模型配置
# 本文中都以 yolox-s 为例
python tools/export_onnx.py --output-name "导出模型名称.onnx" -n yolox-s -c "YOLOX_outputs/yolox_s/best_ckpt.pth"

# 你也可以自定义 Exp
python tools/export_onnx.py --output-name "导出模型名称.onnx" -f "exps/default/yolox_s.py" -c "YOLOX_outputs/yolox_s/best_ckpt.pth"

onnx推理

进入 demo/ONNXRuntime 目录

# -s 指定置信度阈值, 低于该值的结果会被过滤
# --input_shape 指定模型尺寸, 需要与训练时的一致
python onnx_inference.py -m "onnx模型路径.onnx" -i "输入图片" -o "输出目录" -s 0.3 --input_shape 640,640

参考与鸣谢

https://blog.csdn.net/java_pythons/article/details/120513863
https://blog.csdn.net/zxxxiazai/article/details/120906366
https://blog.csdn.net/qq_64693987/article/details/128659749

本文链接：