Training CLI Fork for YOLOv9

[mapo80]

Hi everyone,

I've been working on a fork that adds a complete training infrastructure to the YOLOv9 codebase. The original repository focuses on model architecture and research—this fork adds everything needed to actually train and deploy models.

What's Included

Training Infrastructure

• Complete CLI with YAML configuration and command-line overrides
• PyTorch Lightning integration for clean, scalable training
• Multi-GPU support with automatic DDP
• Mixed precision training (FP16 and BF16)
• Gradient accumulation and clipping

Data Pipeline

• COCO JSON and YOLO TXT format support
• Dataset caching for faster loading
• Augmentation suite: Mosaic 4/9, MixUp, CutMix, RandomPerspective, RandomHSV, RandomFlip

Evaluation & Monitoring

• Full COCO evaluation metrics (mAP, AP50, AP75, AR)
• Confusion matrix and PR curve generation
• Rich progress bars and eval dashboard
• TensorBoard integration

Export

• ONNX export
• TFLite export (FP32, FP16, INT8 quantization)
• SavedModel format

Testing

• 301 unit tests covering data loading, augmentations, metrics, and training pipeline

Repository

https://github.com/mapo80/YOLO

Feedback and contributions are welcome!

[henrytsui000]

Hi,

Thanks so much for your feedback and for sharing your repo—it looks like a great effort and contribution to the ecosystem!

There are two points I’d like to raise:

Feature Overlap: Some of the features you mentioned (like mixed precision training, COCO evaluation, multi-format export, etc.) are already implemented or in progress in this repo. We’re working to keep the core lean while supporting flexible extensions, so any ideas around modularization are appreciated!

License Concerns: I noticed some use of AI agents in code generation. While that’s a valid approach, it raises a concern regarding license compliance. Specifically, it's hard to verify whether AI-generated code has any indirect dependency on non-MIT-licensed implementations. Since this repo is MIT-licensed, I want to ensure all contributions align with that to avoid legal or attribution issues.

Thanks again for sharing, and feel free to open discussions or PRs if you'd like to upstream any modular improvements!

Best,
Henry Tsui

[hitalop]

@mapo80 Hi. After export my model in onnx format, the output shape is [1,10,124,124] should be length 4? And i have 10 classes, the shape should be 4+10. and What does the 124 in the shape represent? where does the component value 124 come from?
applying inference to model.onnx gives a error in postprocess in mask:
IndexError:boolean index did not match indexed array along dimension 1; dimension is 124 but corresponding boolean dimension is 6.

[mapo80]

The output shape [1, 10, 124, 124] indicates that the model was exported without the YOLOv9 export wrapper. This is the raw feature map output from one detection head, not the decoded detection format.

What your shape means:

• 1 = batch size
• 10 = your number of classes (class scores only, no box coordinates!)
• 124, 124 = spatial grid size (this comes from input_size / stride, e.g., 992/8 = 124)

Expected format:

The SDK exports models with a wrapper that:

1. Decodes boxes from anchor offsets to absolute XYWH coordinates
2. Applies sigmoid to class scores
3. Concatenates all detection scales (stride 8, 16, 32)
4. Transposes to [batch, 4+num_classes, num_anchors]

For a 640x640 model with 10 classes, the expected output is:

• [1, 14, 8400] where 14 = 4 (box) + 10 (classes), 8400 = 80×80 + 40×40 + 20×20 anchors

How to fix:

Use the CLI export command:

python -m yolo export --checkpoint your_model.ckpt --format onnx --output model.onnx

This automatically wraps the model with YOLOv9ExportWrapper which produces the correct output format.

Verify your export:

import onnxruntime as ort
session = ort.InferenceSession("model.onnx")
print(f"Input: {session.get_inputs()[0].shape}")
print(f"Output: {session.get_outputs()[0].shape}")

Expected output for 640x640 with 10 classes:

Input: [1, 3, 640, 640]
Output: [1, 14, 8400]

Running inference on ONNX models:

Once exported correctly, you can use the standalone inference script provided at tools/onnx_inference.py:

python tools/onnx_inference.py --model model.onnx --image test.jpg --output result.jpg

This script handles the complete pipeline: letterbox preprocessing, inference, NMS, coordinate scaling, and visualization. It only requires onnxruntime, numpy, and opencv-python (no SDK dependencies).

For more details, see the "Export to ONNX" → "ONNX Model Format Specification" and "ONNX Inference Script" sections in the README.md.

[JanOsusky]

@mapo80 Hello, I ran into an issue while exporting the model to TF.Lite FP16 using the provided Docker image. My finetuned model has 2 classes. It finished training with mAP5.95=0.32.

I exported the model to TF.Lite, which resulted in a model with an output shape of [1, 6, 8400] of one tensor. According to what I read in the [README.md] (https://github.com/mapo80/YOLO/blob/main/README.md), this is the expected shape for a model with 2 classes.

However when I tried to validate my export with the provided validating script, it failed expecting an output of 2 tensors with shapes of [0, 3, 6] and [2, 5, 8] respectively. I assumed that the validation script is incorect.

I tried to quickly modify the script to get some results and I am not sure if it is correct but the model now shows a way worse performence of mAP5.95=0.21. Do you have an idea if I am going in the right direction and what can cause such a significant decrease in accuracy of the model. Thank you.

[mapo80]

Root Cause

The validate_tflite.py script had 80 classes hardcoded (COCO default) in several places:

• Output tensor identification logic
• Class probability reshape operations

When your 2-class model outputs tensors with shape (1, H, W, 2) instead of (1, H, W, 80), the script failed to identify the correct output tensors and decoded the predictions incorrectly — explaining the mAP drop from 0.32 to 0.21.

The Fix

The validate_tflite.py script has been updated to support custom num_classes and both COCO and YOLO dataset formats.

For your case:

# If using COCO-format annotations
python scripts/validate_tflite.py \
    --model your_model_fp16.tflite \
    --ann-file /path/to/annotations.json \
    --images-dir /path/to/images \
    --num-classes 2

# If using YOLO-format dataset (images/ + labels/ directories)
python scripts/validate_tflite.py \
    --model your_model_fp16.tflite \
    --images-dir /path/to/images \
    --labels-dir /path/to/labels \
    --data-format yolo \
    --num-classes 2

Note on TFLite Output

The different output format is expected:

• ONNX: single tensor [1, 6, 8400]
• TFLite: 18 separate tensors (better compatibility)

The export is correct — only the validation script needed the fix.

Summary

With --num-classes 2, mAP values should now be consistent with PyTorch training results (~0.32).

[JanOsusky]

@mapo80 Thank you for your reply and effort. Your fix just doesn't seem to fix my problem. After I used your updated script I ran into the exact same problem:

(venv) root@JANsSlave:/home/school/YOLO# python scripts/validate_tflite.py --model ../exports/model_fp16.tflite --ann-file
/Prices-4/prices4-coco-std/annotations/instances_train2017.json --images-dir ../Prices-4/prices4-coco-std/train2017/ --num-classes 2
Loading model: ../exports/model_fp16.tflite
2026-02-09 00:23:49.583503: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
/home/school/YOLO/venv/lib/python3.10/site-packages/tensorflow/lite/python/interpreter.py:457: UserWarning: Warning: tf.lite.Interpreter is deprecated and is scheduled for deletion in
TF 2.20. Please use the LiteRT interpreter from the ai_edge_litert package.
See the migration guide
for details.
warnings.warn(_INTERPRETER_DELETION_WARNING)
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
Input shape: [ 1 640 640 3]
Input dtype: <class 'numpy.float32'>
Number of outputs: 1
Model type: FP32/FP16
Number of classes: 2
Classification output indices: [0, 3, 6]
Box output indices: [2, 5, 8]
Input size: 640x640
Loading COCO annotations from: ../Prices-4/prices4-coco-std/annotations/instances_train2017.json
Loaded 500 images
Validating: 0%| | 0/500 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/home/school/YOLO/scripts/validate_tflite.py", line 722, in
main()
File "/home/school/YOLO/scripts/validate_tflite.py", line 707, in main
validate_tflite(
File "/home/school/YOLO/scripts/validate_tflite.py", line 591, in validate_tflite
cls_pred, box_pred = decode_outputs(
File "/home/school/YOLO/scripts/validate_tflite.py", line 248, in decode_outputs
box_out = dequantize_output(interpreter, output_details, box_idx)
File "/home/school/YOLO/scripts/validate_tflite.py", line 205, in dequantize_output
det = output_details[index]
IndexError: list index out of range

The script is looking for tesnsor that my model doesn't have. Or thats what I assume it does when it says box and classfications output indices.

What you said about TFLite output

The different output format is expected:

ONNX: single tensor [1, 6, 8400]
TFLite: 18 separate tensors (better compatibility)
The export is correct — only the validation script needed the fix.

For me the TFlite model has the single tensor of [1,6,8400] not 18 tensors. Acoording to what is in the Readme this actually makes sense.

How I understand it

The script probably expects a diffrent output of a modes that still needs some post processing like NMS or similar, but the model that comes out of the conversion already has those applied. Am I right? Do you know how to fix this error?

Thank you

Edit: just what AI told me

Looking at the error and your description, the issue is clear: your model outputs a single tensor of shape [1, 6, 8400] (the classic "concatenated" YOLO export format), but the validation script was written expecting the multi-head split output format with 18 separate tensors.
This is a completely different output format — your model has already concatenated everything into one tensor where the 6 channels are [x, y, w, h, class0_score, class1_score] (or similar), flattened across 8400 anchors.

I don't know if the converting script to .tflite is expected to make a not processed prediction or some kind of already finished detection that already directly returns just bbox and class and confidence or make those 18 tensors as you said, either way its clearly not working the way the validation scrip expect its to work

[JanOsusky]

Hello @mapo80, I would like to bring attention back to this matter. Thank you

[Njaecha]

This fork has been very useful tome since it could actually do the training, while I couldn't get the training in this repository to work.
That and the actually documented onnx export.
However running it on windows, the caching feature can't be used since the LMDB causes problems on multiple fronts.
Since I didn't want to spend a lot of time trying to fix the caching when I could also just wait longer for the training, I had Claude take a look at the caching and after a few rounds of back and forth it kinda just made it worse.

I would raise and issue with more details on the fork itself, but there issues seem to be disabled.