记一次训练 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