Exclude Sigmoid from TensorRT INT8 quantization - #25020

All Contributors have signed the CLA. ✅
_{Posted by the CLA Assistant Lite bot.}

👋 Hello @davidnichols-ops, thank you for submitting a ultralytics/ultralytics 🚀 PR! This automated message is here to help with the initial review process, and an Ultralytics engineer will assist further. Your PR appears to address YOLO26 INT8 confidence calibration, TensorRT export behavior, end2end=False score handling, tests, and related documentation updates.

• ✅ Define a Purpose: Clearly explain the purpose of your fix or feature in your PR description, and link to any relevant issues. Ensure your commit messages are clear, concise, and adhere to the project's conventions.
• ✅ Synchronize with Source: Confirm your PR is synchronized with the ultralytics/ultralytics main branch. If it's behind, update it by clicking the 'Update branch' button or by running git pull and git merge main locally.
• ✅ Ensure CI Checks Pass: Verify all Ultralytics Continuous Integration (CI) checks are passing. If any checks fail, please address the issues.
• ✅ Update Documentation: Update the relevant documentation for any new or modified features.
• ✅ Add Tests: If applicable, include or update tests to cover your changes, and confirm that all tests are passing.
• ✅ Sign the CLA: Please ensure you have signed our Contributor License Agreement if this is your first Ultralytics PR by writing "I have read the CLA Document and I sign the CLA" in a new message.
• ✅ Minimize Changes: Limit your changes to the minimum necessary for your bug fix or feature addition. "It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is." — Bruce Lee

For more guidance, please refer to our Contributing Guide. Don't hesitate to leave a comment if you have any questions. Thank you for contributing to Ultralytics! 🚀

Codecov Report

❌ Patch coverage is 7.69231% with 12 lines in your changes missing coverage. Please review.

📢 Thoughts on this report? Let us know!

@Y-T-G Yes — the PR now contains only Fix 2 and Fix 3. Fix 1 (exporter warning for end2end=False) and Fix 4 (one-to-one head selection for end2end=False) have been removed entirely.

The squashed commit (5cfc713) touches only:

• ultralytics/utils/export/engine.py — Sigmoid/Softmax exclusion from INT8 quantization (TRT 10 per-layer FP32 constraint + TRT 11+ ModelOpt op_types_to_exclude)
• tests/test_exports.py — tests for the above
• docs/en/integrations/tensorrt.md — short doc note
• tests/bench_int8_retention.py — benchmark script

No changes to head.py or exporter.py's end2end logic remain. The stale inline comments from UltralyticsAssistant on those files have been resolved.

Latest changes — what I found and fixed

While testing the Softmax generalization, I ran into a few issues with the TRT 10 path and the test suite. Here's what happened and how I dealt with each:

1. Softmax wasn't actually being preserved on TRT 10 (real bug)

In TensorRT, Softmax is a separate ISoftMaxLayer with LayerType.SOFTMAX — it's not an IActivationLayer. My original loop only checked layer.type == trt.LayerType.ACTIVATION, so Softmax layers were filtered out before the detection helper ever ran. trt.ActivationType also has no SOFTMAX member, so the typed-cast path was dead code for Softmax too. The Sigmoid fix worked fine (Sigmoid is an activation layer), but the Softmax generalization was advertised without actually being implemented.

Fixed by adding an explicit LayerType.SOFTMAX branch in _constrain_sensitive_activations_fp32:

softmax_type = getattr(trt.LayerType, "SOFTMAX", None)
...
if layer.type == trt.LayerType.ACTIVATION and _is_precision_sensitive_activation(layer, trt):
    # constrain Sigmoid activations
elif softmax_type is not None and layer.type == softmax_type:
    # constrain Softmax layers

The execution test now includes a Softmax layer with LayerType.SOFTMAX and asserts it's constrained (count == 2: Sigmoid + Softmax). For the reported YOLO26 detect issue only Sigmoid matters, so the fix was already working for its stated purpose — but the Softmax generalization is now actually implemented rather than just advertised.

2. The call-site test wasn't actually testing the call site

The old test_trt10_sigmoid_fp32_precision_constraint just called the helper directly with an empty network and checked the return type. It didn't verify that onnx2engine actually invokes the helper at the right time. Replaced it with test_trt10_constrain_sensitive_activations_call_site that builds a full fake tensorrt module (Builder, Config, Network, Parser, Logger, IInt8Calibrator, all flags/enums), creates a minimal valid ONNX file, and calls onnx2engine directly. It verifies the helper is:

• Called for INT8 on TRT 10
• NOT called for FP16 on TRT 10
• NOT called for INT8 on TRT 11 (ModelOpt handles it there)

3. Comment said FP32, but it's actually FP16

The comment on the ModelOpt path said excluded nonlinearities "run in FP32 on dequantized logits," but ModelOpt defaults to high_precision_dtype='fp16', so they actually run in FP16. Fixed the comment to say so, and noted that OpenVINO's IgnoredScope is stricter (keeps ops in FP32) but the accuracy impact is negligible for 0-3 logits.

4. Test could trigger a network install

The ModelOpt test was calling check_requirements("nvidia-modelopt[onnx]>=0.44") which, if modelopt is absent and autoinstall is on, would try to pip install it. Safe in CI (where YOLO_AUTOINSTALL=false), but not hermetic. Added a monkeypatch for check_requirements so the test is self-contained.

5. Removed a test that's no longer relevant

Dropped test_end2end_false_export_smoke — it tested end2end=False export behavior, which was part of the original PR but was removed per the earlier review feedback to scope down to Fix 2 and Fix 3 only.

6. OpenVINO parity test — kept as-is

The OpenVINO test uses source inspection (searching the wrapped function's source for IgnoredScope and Sigmoid). It's not as strong as a behavior test, but a behavior-driven test would require installing OpenVINO + NNCF, which is too heavy for CI. The closure-cell iteration was improved to search for the first callable cell rather than indexing a fixed slot, so it survives ordering changes if try_export is refactored.

Verification:

$ ruff check ultralytics/utils/export/engine.py tests/test_exports.py
All checks passed!
$ pytest tests/test_exports.py — 27 passed, 5 failed (pre-existing missing optional deps), 4 skipped

Thanks, this rescope addresses the original concern and the TRT 10 Softmax branch/test coverage makes sense. One remaining cleanup to keep the PR minimal: please drop the standalone tests/bench_int8_retention.py benchmark script from the committed diff and keep the benchmark/reproduction details in the PR description or docs instead; the focused engine.py, test_exports.py, and TensorRT docs changes should be enough.

Done — dropped tests/bench_int8_retention.py from the diff. The committed changes are now just the three files:

• ultralytics/utils/export/engine.py — Sigmoid/Softmax exclusion logic
• tests/test_exports.py — tests for the above
• docs/en/integrations/tensorrt.md — doc note

Benchmark/reproduction details remain in the PR description (collapsed sections) and the Kaggle kernel link.

Thanks, that cleanup is exactly what I was looking for. With the benchmark script removed, the PR is now focused on the TensorRT INT8 precision handling plus tests/docs; we can proceed with the normal CI/review path from here.

Thanks, Glenn. I appreciate the feedback. I’ll leave it as-is and wait for CI approval and the normal review process!

I have updated the PR to include the requested type annotations in _constrain_sensitive_activations_fp32 to resolve the documentation CI gap.

Thanks, the type annotations on _constrain_sensitive_activations_fp32 are the right cleanup for the docs CI issue. Let’s let CI rerun on the updated diff and continue through the normal review path.

Thanks @davidnichols-ops, the diagnosis and T4 validation here are excellent and the approach (mirroring OpenVINO's IgnoredScope) is exactly right. 🙏

I've pushed a simplification to your branch that preserves the validated behavior while reducing the diff to +35/−4:

• The trt.IActivationLayer(layer) typed cast isn't constructible in the TRT 10 Python bindings (it raises TypeError on every version, not just 10.7+), so the name-matching fallback was the code path that actually ran in your T4 benchmarks. The revision keeps just that path: Softmax by LayerType.SOFTMAX, Sigmoid by ONNX node name — folded into the existing use_int8 and not is_trt11 calibrator branch.
• op_types_to_exclude is now inlined at both call sites instead of the paired module constants.
• Kept one focused unit test guarding the ModelOpt kwarg; your hardware validation covers the TRT 10 loop end-to-end.
• Trimmed the docs to a short tip (docs shouldn't point users at private constants).

Should this be quantize=w8a32 just like LiteRT?

Good question — I looked into this, and I don’t think w8a32 applies directly to this particular path.

From what I found, ModelOpt’s ONNX quantization flow does not currently expose a weight-only mode. The available modes are INT8 static, INT4, and FP8. For the INT8 path this PR is using, excluded ops are kept in high_precision_dtype (default FP16), so op_types_to_exclude appears to be the intended mechanism for keeping numerically sensitive ops like Sigmoid/Softmax out of INT8.

The TRT 10 path also seems limited here since it uses activation + weight INT8 through IInt8Calibrator; I couldn’t find a weight-only builder option for this flow. LiteRT can provide w8a32 because its quantizer has a separate weight-only recipe, but I don’t see an equivalent exposed through ModelOpt/TRT ONNX currently.

I also ran a quick T4 comparison using ONNX Runtime quantize_dynamic (INT8 weights + FP32 activations) as a w8a32 proxy alongside the current TRT INT8 path (ModelOpt quantize=8 with Sigmoid/Softmax exclusion). The weight-only proxy did preserve confidence values better in my test, which is expected since activations remain higher precision.

Since ModelOpt doesn’t currently expose a weight-only mode, would it be worth considering whether an ONNX Runtime-based weight-only path is something worth exploring, or would the added complexity outweigh the accuracy benefit compared with the current ModelOpt/TensorRT approach?

I can share the benchmark details if useful.

Conclusion for this PR: I would keep the existing quantize=8 path with Sigmoid/Softmax exclusion rather than switching to w8a32, since ModelOpt/TensorRT does not currently provide an equivalent weight-only mode.

Agreed — let’s keep this PR on the current TensorRT quantize=8 path with Sigmoid/Softmax exclusion. w8a32 is currently a LiteRT-specific dynamic/weight-only export mode in our export precision matrix, while TensorRT INT8 export remains calibration-based as described in the TensorRT docs. An ONNX Runtime weight-only route would be a separate feature/design PR if we decide it’s worth supporting, not something to add here.

Thank you Glenn for the help in getting this to the finish line. Nothing like green CI's in realtime.

Thanks @davidnichols-ops — your careful diagnosis, T4 validation, and responsive cleanup made this a strong community contribution. Glad we got it merged with green CI!

Thanks again @davidnichols-ops—the careful validation and focused iteration benefited the whole YOLO community. We appreciate the contribution and look forward to your future PRs.