Ternary QAT YOLO26s drill rig detector
Small-object detection with ternary QAT and a custom C++ inference runtime
Metrics
- 15,311 labeled images, single-class drill rig detection
- 43 targeted 1x1 convolutions ternarized across the YOLO26s backbone, neck, and attention blocks
- Per-layer sensitivity profiling, mAP50-95 0.95216 to 0.95207 at one ternary layer, 0.91056 at ten, decided which layers were safe
- ~0.2-0.25 bytes per ternary weight via 5-trit packing, vs. 4 bytes in FP32
- Custom AVX2 + OpenMP C++ ONNX Runtime operator, ~3.28s to ~0.78s per inference, a ~4x speedup
Tech stack
- Python
- PyTorch
- Ultralytics YOLO
- C++17
- ONNX Runtime
- CMake
What it does
This detector finds drill rigs in imagery where the rig itself is a small object relative to the frame, the kind of case where general-purpose object detectors tend to lose recall. It's built on YOLO26s and trained on 15,311 labeled images, but the project is really two pieces: compressing the model down to ternary weights, and then writing an inference engine that can actually run on them.
The compression side is quantization-aware training down to ternary values, 1, applied to 43 targeted 1x1 convolutions across the backbone, neck, and attention blocks. Each of those convolutions keeps a trainable full-precision shadow weight that gets projected per output-channel to ternary via AbsMean scaling, trained through the discretization with a straight-through gradient estimator. Rather than ternarizing everything at once, I profiled the accuracy impact layer by layer first: mAP50-95 barely moves, 0.95216 to 0.95207, with a single ternary layer, but drops to 0.91056 by ten, and that curve is what decided which of the 43 layers could absorb the precision loss without hurting detection quality.
Ternary weights only help if something downstream can use them without unpacking back to full precision. I packed them to roughly 0.2-0.25 bytes per weight, five trits per byte for storage, 2-bit codes at runtime, versus 4 bytes per weight in FP32, then wrote a custom ONNX Runtime operator in C++ that runs a native ternary 1x1 convolution kernel directly against that packed format, with AVX2 vectorization, OpenMP multithreading, and paired-channel execution to reuse activation reads. It's registered as a custom op and exercised through a full ONNX graph, with a CMake build and layer-level parity tests that pass on both Windows and Linux.
The runtime went from an initial ~3.28s per inference down to ~0.78s, a roughly 4x speedup, through packed storage, sparser execution, and AVX2/OpenMP tuning. It's still about 3x slower than PyTorch's own CPU inference on the same hardware, which is the main open problem right now. Activations are still FP32, so this is weight-only compression for the moment, not full ternary x INT8 inference, that's the planned next step. The work is ongoing, with an IEEE paper in preparation.
