Adaptation of the YOLOX framework (Megvii BaseDetection) to train a custom detector for faces (face) and identity documents (doc_quad), and export it to ONNX for downstream inference (browser, edge, server — anywhere with an ONNX runtime).
Fine-tune YOLOX-S on a custom VOC dataset (derived from MIDV-2020) with 2 classes (
face,doc_quad). Config inexps/example/custom/yolox_doc_face.py. Train withtools/train.py, export to ONNX withexport_correct_ONNX.py(which enablesdecode_in_inferenceso the consumer doesn't have to decode boxes), and validate withcompare_models.py.
This repo covers only training and ONNX export. The resulting .onnx is a self-contained artifact you can plug into any ONNX-compatible runtime (onnxruntime, onnxruntime-web, TensorRT, OpenVINO, etc.). Real-time inference clients live outside this repo.
End-to-end flow:
┌───────────────┐ ┌────────────┐ ┌───────────────┐ ┌──────────────┐
│ MIDV-2020 │ → │ Convert to │ → │ Training │ → │ ONNX Export │
│ (download) │ │ VOC XML │ │ (fine-tune S) │ │ (.pth → .onnx)│
└───────────────┘ └────────────┘ └───────────────┘ └──────────────┘
YOLOX is an anchor-free object detector from the YOLO family (Megvii, 2021). Key features:
- Anchor-free — no predefined anchor boxes, simpler design.
- Decoupled head — separates classification and regression.
- SimOTA — dynamic label assignment during training.
- Strong augmentation — Mosaic + MixUp by default.
- Variants:
Nano,Tiny,S,M,L,X,Darknet53.
This project starts from the COCO-pretrained YOLOX-S checkpoint (yolox_s.pth) and fine-tunes it.
yolox-trainer/
├── yolox/ # Framework (models, data, utils, evaluators)
│ ├── core/ # Trainer, launcher
│ ├── data/ # Datasets (COCO, VOC), dataloaders, augment
│ ├── evaluators/ # COCO / VOC evaluators
│ ├── exp/ # Base Exp class
│ ├── layers/ # Custom layers
│ ├── models/ # Backbone, FPN, YOLOXHead
│ ├── tools/ # Module entrypoints
│ └── utils/ # Logger, metrics, postprocess
│
├── tools/ # CLI scripts
│ ├── train.py # Training
│ ├── eval.py # Evaluation
│ ├── demo.py # Image/video inference
│ ├── export_onnx.py # Generic ONNX export (upstream)
│ └── trt.py # TensorRT export
│
├── exps/
│ ├── default/ # Official configs (yolox_s, yolox_m, …)
│ └── example/
│ ├── custom/
│ │ └── yolox_doc_face.py # ⭐ Custom config
│ └── yolox_voc/ # VOC reference templates
│
├── datasets/ # Data (NOT included in repo!)
│ └── VOCdevkit/VOC2020/
│
├── demo/ # Demos per deployment backend
│ ├── ONNXRuntime/
│ └── TensorRT/, ncnn/, OpenVINO/, MegEngine/, nebullvm/
│
├── docs/ # Upstream framework tutorials
│
├── compare_models.py # ⭐ Compare PyTorch vs ONNX output
├── export_correct_ONNX.py # ⭐ ONNX export tuned for this model
├── hubconf.py # torch.hub entrypoint
├── requirements.txt
├── setup.py / setup.cfg
└── README.md # ← this document
⭐ = added/customized for this project.
exps/example/custom/yolox_doc_face.py:
class Exp(MyExp):
def __init__(self):
self.num_classes = 2
self.depth = 0.33 # "S" size
self.width = 0.50
self.warmup_epochs = 1
self.class_names = ["doc_quad", "face"]
self.mosaic_prob = 1.0
self.mixup_prob = 1.0
self.hsv_prob = 1.0
self.flip_prob = 0.5
self.exp_name = "yolox_doc_face"- Classes:
doc_quad(idx 0),face(idx 1). - Backbone: YOLOX-S (
depth=0.33,width=0.50). - Dataset:
VOCDetectionfromdatasets/VOCdevkit/with splits("2020", "train")and("2020", "val"). - Evaluator:
VOCEvaluator→ Pascal VOC-style mAP metrics.
The repo does not ship the data. You must download and convert it to VOC manually.
- MIDV-2020: identity document images.
- Official site: http://l3i-share.univ-lr.fr/MIDV2020/midv2020.html
datasets/VOCdevkit/
├── annotations_cache/ # Auto-generated (DELETE if you add data)
└── VOC2020/
├── Annotations/ # *.xml (VOC format)
├── JPEGImages/ # *.jpg
└── ImageSets/Main/
├── train.txt # filenames without extension
└── val.txt
<annotation>
<filename>image001.jpg</filename>
<size><width>2268</width><height>4032</height><depth>3</depth></size>
<object>
<name>face</name>
<bndbox><xmin>100</xmin><ymin>200</ymin><xmax>300</xmax><ymax>400</ymax></bndbox>
</object>
<object>
<name>doc_quad</name>
<bndbox><xmin>500</xmin><ymin>600</ymin><xmax>800</xmax><ymax>900</ymax></bndbox>
</object>
</annotation>- Each
image001.jpgrequiresimage001.xmlwith the same base name. train.txt/val.txt: one filename per line, without extension.- Typical 80/20 split.
⚠️ If you add new images, deletedatasets/VOCdevkit/annotations_cache/before the next training run.
- Python 3.8 (Anaconda recommended)
- CUDA 11.8 + NVIDIA GPU (recommended; CPU works but is very slow)
- Windows 10/11, macOS or Linux
# 1. Environment
conda create -n yolox python=3.8 -y
conda activate yolox
# 2. PyTorch (adjust CUDA to your system)
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -y
# CPU-only: conda install pytorch torchvision torchaudio cpuonly -c pytorch -y
# 3. YOLOX in editable mode
cd /path/to/yolox-trainer
pip install -v -e .
# 4. Extra dependencies
pip install cython pycocotools tensorboard onnx onnxruntime
# 5. Verify
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
python -c "import yolox; print('OK')"
# 6. Base YOLOX-S weights
# Download manually from:
# https://github.com/Megvii-BaseDetection/YOLOX/releases
# and place yolox_s.pth at the repo root.conda activate yolox
python tools/train.py \
-f exps/example/custom/yolox_doc_face.py \
-d 1 \ # 1 GPU
-b 8 \ # total batch size
--fp16 \ # mixed precision
-o \ # occupy GPU memory upfront
-c yolox_s.pth # initial checkpoint (transfer learning)Useful flags:
-d N: number of GPUs.-b N: total batch size. Recommended8 × num_gpus. If OOM, lower it.--cache: cache images in RAM (faster, requires lots of RAM).--logger wandb: integrate with Weights & Biases.
YOLOX_outputs/yolox_doc_face/
├── train_log.txt
├── tensorboard logs
├── latest_ckpt.pth
├── best_ckpt.pth ← exported to ONNX
└── epoch_*.pth
python tools/eval.py \
-f exps/example/custom/yolox_doc_face.py \
-c YOLOX_outputs/yolox_doc_face/best_ckpt.pth \
-b 8 -d 1 --conf 0.001 --fp16 --fuseReported metrics: Pascal VOC-style mAP (not COCO), because the experiment uses VOCEvaluator.
Use export_correct_ONNX.py (tuned for this model). Do not use the generic tools/export_onnx.py — it doesn't set critical params like decode_in_inference=True.
python export_correct_ONNX.py- Loads
exps/example/custom/yolox_doc_face.py. - Sets
num_classes=2,test_conf=0.25,nmsthre=0.45,test_size=(640, 640). - Loads the checkpoint (
⚠️ path is hardcoded in the script — update it if it changes). - Enables
model.head.decode_in_inference = True→ box decoding happens inside the ONNX graph. - Exports with
opset_version=11,dynamic_axesfor dynamic batch,do_constant_folding=True. - Verifies the ONNX by loading it with
onnxruntimeand comparing outputs against PyTorch (mean diff must be < 0.001).
yolox_doc_face_fixed.onnxat the repo root.- Output tensor:
[batch, num_anchors, 7]containing[x, y, w, h, objectness, p_doc_quad, p_face]already decoded.
python compare_models.py \
--image prueba_doc.jpg \
--checkpoint YOLOX_outputs/yolox_doc_face_20250805/best_ckpt.pth \
--onnx yolox_doc_face_fixed.onnxPrints PyTorch detections (postprocess + NMS) and the top-3 ONNX candidates per class. Useful to catch regressions after re-exporting.
python tools/demo.py image \
-f exps/example/custom/yolox_doc_face.py \
-c YOLOX_outputs/yolox_doc_face/best_ckpt.pth \
--path assets/dog.jpg \
--conf 0.25 --nms 0.45 --tsize 640 \
--save_result --device gpupython tools/demo.py video \
-f exps/example/custom/yolox_doc_face.py \
-c YOLOX_outputs/yolox_doc_face/best_ckpt.pth \
--path /path/to/video.mp4 \
--conf 0.25 --nms 0.45 --tsize 640 \
--save_result --device gpuResults land in YOLOX_outputs/yolox_doc_face/vis_res/.
python demo/ONNXRuntime/onnx_inference.py \
-m yolox_doc_face_fixed.onnx \
-i prueba_doc.jpg| Symptom | Likely cause | Fix |
|---|---|---|
CUDA out of memory |
Batch too large / small GPU | Lower -b (e.g. -b 4 or -b 2) |
module 'yolox' not found |
Env not activated or broken install | conda activate yolox and pip install -v -e . |
| ONNX detections in wrong locations | Forgot decode_in_inference=True |
Re-export with export_correct_ONNX.py |
| Weird metrics after adding images | Stale cache | rm -rf datasets/VOCdevkit/annotations_cache/ |
onnxruntime can't find the model |
Wrong relative path | Use absolute path or --onnx <path> |
| PyTorch vs ONNX diff > 0.001 | Incompatible opset or unexported layers | Change opset_version or check custom layers |
- YOLOX paper: https://arxiv.org/abs/2107.08430
- Upstream repo: https://github.com/Megvii-BaseDetection/YOLOX
- MIDV-2020 dataset: http://l3i-share.univ-lr.fr/MIDV2020/midv2020.html
- Framework tutorials (
docs/folder):train_custom_data.md— training on custom datacache.md— cache optionsfreeze_module.md— freezing modulesassignment_visualization.md— SimOTA visualization
If you use this work in research, please cite the original YOLOX paper:
@article{yolox2021,
title={YOLOX: Exceeding YOLO Series in 2021},
author={Ge, Zheng and Liu, Songtao and Wang, Feng and Li, Zeming and Sun, Jian},
journal={arXiv preprint arXiv:2107.08430},
year={2021}
}Apache License 2.0 — inherited from the upstream YOLOX project. See LICENSE.