Fix TensorRT GridSample inference for LT-DETR

CHANGELOG.md

@@ -22,6 +22,13 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Fixed
 
+- Fix incorrect TensorRT inference for LT-DETR object detection and instance
+  segmentation. TensorRT's optimizer/fusion pass around the `GridSample` ops used by
+  deformable attention silently produced wrong activations, corrupting detections and
+  masks, even with parser-compatible mode names. Deployment/export now replaces
+  `grid_sample` with a gather-based bilinear equivalent that contains no `GridSample`
+  op, so neither the parser mode-name issue nor the optimizer bug can apply; training
+  keeps the faster fused `grid_sample`.
 - Fix `export_onnx(verify=True)` crashing with an `onnx.checker` `ShapeInferenceError`
   for LTDETR object detection by repairing stale intermediate shape annotations left by
   the ONNX exporter.

src/lightly_train/_export/export_onnx.py

@@ -24,6 +24,7 @@
 from lightly_train._export.export import ExportMixin
 from lightly_train._export.onnx_helpers import (
     fix_topological_order,
+    prepare_for_onnx_export,
     remove_duplicate_cast_nodes,
     remove_redundant_casts,
     repair_value_info,
@@ -111,6 +112,7 @@ def export_onnx(
         # safely after export. The shared pipeline deliberately preserves that behavior.
         module.to(torch.float32)
         module.deploy()
+        prepare_for_onnx_export(module)
 
         program = self.export(
             batch_size=batch_size,

src/lightly_train/_export/onnx_helpers.py

@@ -13,10 +13,11 @@
 from collections.abc import Iterator
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 import torch
 from lightning_utilities.core.imports import RequirementCache
+from torch.nn import Module
 
 if TYPE_CHECKING:
     import onnx
@@ -30,6 +31,11 @@
 _TORCH_DIM_HINTS_AVAILABLE = RequirementCache(f"torch>={_TORCH_DIM_HINTS_MIN_VERSION}")
 
 
+@runtime_checkable
+class ONNXExportConvertible(Protocol):
+    def convert_to_onnx_export(self) -> None: ...
+
+
 def check_onnx_dynamo_requirements() -> None:
     """Raise if the installed torch version does not support dynamo ONNX export."""
     if not _TORCH_DYNAMO_AVAILABLE:
@@ -113,6 +119,13 @@ def repair_value_info(out: str | Path) -> None:
     onnx.save(model, str(out))
 
 
+def prepare_for_onnx_export(module: Module) -> None:
+    """Apply module-specific graph conversions required for ONNX export."""
+    for child in module.modules():
+        if isinstance(child, ONNXExportConvertible):
+            child.convert_to_onnx_export()
+
+
 def remove_duplicate_cast_nodes(model: onnx.ModelProto) -> None:
     """Remove duplicate Cast nodes emitted by convert_float_to_float16 in-place.
 

src/lightly_train/_export/tensorrt_helpers.py

@@ -155,10 +155,11 @@ def export_tensorrt(
 
     logger.info(f"Loading ONNX file from {onnx_out}")
     with open(onnx_out, "rb") as f:
-        if not parser.parse(f.read()):
-            for error in range(parser.num_errors):
-                logger.error(parser.get_error(error))
-            raise RuntimeError("Failed to parse ONNX file")
+        onnx_bytes = f.read()
+    if not parser.parse(onnx_bytes):
+        for error in range(parser.num_errors):
+            logger.error(parser.get_error(error))
+        raise RuntimeError("Failed to parse ONNX file")
 
     if fp32_attention_scores:
         _force_fp32_for_attention_scores(network)

src/lightly_train/_task_models/ltdetr_instance_segmentation/task_model.py

@@ -26,6 +26,7 @@
 from lightly_train._export import tensorrt_helpers
 from lightly_train._export.onnx_helpers import (
     fix_topological_order,
+    prepare_for_onnx_export,
     remove_redundant_casts,
 )
 from lightly_train._models import package_helpers
@@ -374,6 +375,7 @@ def export_onnx(
         # post-export via onnxruntime.transformers.
         self.to(torch.float32)
         self.deploy()
+        prepare_for_onnx_export(self)
         model_device = next(self.parameters()).device
 
         # Infer num_channels if not provided. The model always consumes

src/lightly_train/_task_models/object_detection_components/dfine_decoder.py

@@ -158,6 +158,18 @@ def _reset_parameters(self):
         init.constant_(self.attention_weights.weight, 0)
         init.constant_(self.attention_weights.bias, 0)
 
+    def convert_to_onnx_export(self) -> None:
+        # Switch the default (grid_sample) sampler to a gather-based bilinear
+        # equivalent that exports to a TensorRT-compatible graph. TensorRT
+        # converts ONNX GridSample incorrectly for this deformable attention, so
+        # the fused grid_sample is kept only for training. Discrete sampling is
+        # left unchanged.
+        if self.method == "default":
+            self.method = "bilinear_gather"
+            self.ms_deformable_attn_core = functools.partial(
+                deformable_attention_core_func_v2, method=self.method
+            )
+
     def forward(
         self,
         query,

src/lightly_train/_task_models/object_detection_components/rtdetrv2_decoder.py

@@ -143,6 +143,18 @@ def _reset_parameters(self):
         init.xavier_uniform_(self.output_proj.weight)
         init.constant_(self.output_proj.bias, 0)
 
+    def convert_to_onnx_export(self) -> None:
+        # Switch the default (grid_sample) sampler to a gather-based bilinear
+        # equivalent that exports to a TensorRT-compatible graph. TensorRT
+        # converts ONNX GridSample incorrectly for this deformable attention, so
+        # the fused grid_sample is kept only for training. Discrete sampling is
+        # left unchanged.
+        if self.method == "default":
+            self.method = "bilinear_gather"
+            self.ms_deformable_attn_core = functools.partial(
+                deformable_attention_core_func_v2, method=self.method
+            )
+
     def forward(
         self,
         query: torch.Tensor,

src/lightly_train/_task_models/object_detection_components/utils.py

@@ -37,6 +37,80 @@ def bias_init_with_prob(prior_prob=0.01):
     return bias_init
 
 
+def bilinear_grid_sample(
+    im: Tensor, grid: Tensor, align_corners: bool = False
+) -> Tensor:
+    """Pure gather-based equivalent of ``F.grid_sample`` (bilinear, zero padding).
+
+    ``torch.nn.functional.grid_sample`` exports to an ONNX ``GridSample`` op that
+    TensorRT converts incorrectly for the deformable-attention sampling used by
+    LT-DETR (the engine produces wrong outputs even though the inputs are exact).
+    This decomposition into ``Pad`` + ``Gather`` + arithmetic matches
+    ``F.grid_sample(mode="bilinear", padding_mode="zeros", align_corners=...)`` to
+    within floating-point tolerance and builds a correct TensorRT engine.
+
+    Args:
+        im:
+            Input feature map of shape ``(N, C, H, W)``.
+        grid:
+            Sampling grid of shape ``(N, H_out, W_out, 2)`` with coordinates
+            normalized to ``[-1, 1]``.
+        align_corners:
+            Matches the ``align_corners`` argument of ``F.grid_sample``.
+
+    Returns:
+        Sampled tensor of shape ``(N, C, H_out, W_out)``.
+    """
+    n, c, h, w = im.shape
+    _, gh, gw, _ = grid.shape
+
+    x = grid[..., 0]
+    y = grid[..., 1]
+    if align_corners:
+        x = ((x + 1) / 2) * (w - 1)
+        y = ((y + 1) / 2) * (h - 1)
+    else:
+        x = ((x + 1) * w - 1) / 2
+        y = ((y + 1) * h - 1) / 2
+
+    x = x.reshape(n, -1)
+    y = y.reshape(n, -1)
+    x0 = torch.floor(x)
+    y0 = torch.floor(y)
+    x1 = x0 + 1
+    y1 = y0 + 1
+
+    # Bilinear interpolation weights (computed before clamping so out-of-bounds
+    # samples keep their correct weights and only fetch zeros).
+    wa = ((x1 - x) * (y1 - y)).unsqueeze(1)
+    wb = ((x1 - x) * (y - y0)).unsqueeze(1)
+    wc = ((x - x0) * (y1 - y)).unsqueeze(1)
+    wd = ((x - x0) * (y - y0)).unsqueeze(1)
+
+    # Zero-pad by one pixel on each side so that out-of-bounds coordinates land on
+    # the zero border, reproducing ``padding_mode="zeros"``.
+    im = F.pad(im, [1, 1, 1, 1])
+    padded_h = h + 2
+    padded_w = w + 2
+    x0 = (x0 + 1).long().clamp(0, padded_w - 1)
+    x1 = (x1 + 1).long().clamp(0, padded_w - 1)
+    y0 = (y0 + 1).long().clamp(0, padded_h - 1)
+    y1 = (y1 + 1).long().clamp(0, padded_h - 1)
+
+    im = im.reshape(n, c, padded_h * padded_w)
+
+    def _gather(xx: Tensor, yy: Tensor) -> Tensor:
+        idx = (xx + yy * padded_w).unsqueeze(1).expand(-1, c, -1)
+        return torch.gather(im, 2, idx)
+
+    ia = _gather(x0, y0)
+    ib = _gather(x0, y1)
+    ic = _gather(x1, y0)
+    id_ = _gather(x1, y1)
+
+    return (ia * wa + ib * wb + ic * wc + id_ * wd).reshape(n, c, gh, gw)
+
+
 def deformable_attention_core_func(
     value, value_spatial_shapes, sampling_locations, attention_weights
 ):
@@ -115,7 +189,7 @@ def deformable_attention_core_func_v2(
     value_list = value.permute(0, 2, 3, 1).flatten(0, 1).split(split_shape, dim=-1)
 
     # sampling_offsets [8, 480, 8, 12, 2]
-    if method == "default":
+    if method in ("default", "bilinear_gather"):
         sampling_grids = 2 * sampling_locations - 1
 
     elif method == "discrete":
@@ -138,6 +212,14 @@ def deformable_attention_core_func_v2(
                 align_corners=False,
             )
 
+        elif method == "bilinear_gather":
+            # Gather-based equivalent of the "default" grid_sample that exports to
+            # a TensorRT-compatible graph (see ``bilinear_grid_sample``). Used for
+            # deployment/export; training keeps the faster fused grid_sample.
+            sampling_value_l = bilinear_grid_sample(
+                value_l, sampling_grid_l, align_corners=False
+            )
+
         elif method == "discrete":
             # n * m, seq, n, 2
             sampling_coord = (

tests/_task_models/ltdetr_instance_segmentation/test_task_model.py

@@ -300,6 +300,12 @@ def test_export_onnx__dynamic_batch_size(
     input_batch_dim = onnx_model.graph.input[0].type.tensor_type.shape.dim[0]
     assert input_batch_dim.dim_param == "N"
 
+    # The graph must not contain a GridSample op: TensorRT converts it incorrectly
+    # for the deformable-attention sampling shared with object detection, so
+    # deployment swaps it for a gather-based bilinear equivalent.
+    op_types = {node.op_type for node in onnx_model.graph.node}
+    assert "GridSample" not in op_types
+
     # Use a batch size (3) different from the one used during tracing (2).
     inputs = np.random.randn(3, 3, 256, 256).astype(np.float32)
     orig_target_size = np.array([[256, 256]] * 3, dtype=np.int64)

tests/_task_models/ltdetr_object_detection/test_task_model.py

@@ -842,7 +842,9 @@ def test_export_onnx__dynamic_batch_size(tmp_path: Path) -> None:
 )
 def test_export_onnx__checks(tmp_path: Path) -> None:
     # The graph must pass the full ONNX checker, which requires the dynamo
-    # shape-annotation repair.
+    # shape-annotation repair. It must also not contain a GridSample op: TensorRT
+    # converts it incorrectly for the deformable-attention sampling, so deployment
+    # swaps it for a gather-based bilinear equivalent.
     import onnx
 
     model = LTDETRObjectDetection(
@@ -858,6 +860,10 @@ def test_export_onnx__checks(tmp_path: Path) -> None:
 
     onnx.checker.check_model(str(out), full_check=True)
 
+    onnx_model = onnx.load(out)
+    op_types = {node.op_type for node in onnx_model.graph.node}
+    assert "GridSample" not in op_types
+
 
 @pytest.mark.skipif(not RequirementCache("onnx"), reason="onnx not installed")
 @pytest.mark.skipif(

tests/_task_models/object_detection_components/test_utils.py

@@ -7,9 +7,15 @@
 #
 from __future__ import annotations
 
+import pytest
 import torch
+import torch.nn.functional as F
 
-from lightly_train._task_models.object_detection_components.utils import _yolo_to_xyxy
+from lightly_train._task_models.object_detection_components.utils import (
+    _yolo_to_xyxy,
+    bilinear_grid_sample,
+    deformable_attention_core_func_v2,
+)
 
 
 def test_yolo_to_xyxy_accepts_1d_box() -> None:
@@ -52,3 +58,59 @@ def test_yolo_to_xyxy_accepts_two_boxes() -> None:
         dtype=torch.float32,
     )
     torch.testing.assert_close(converted[0], expected)
+
+
+@pytest.mark.parametrize(
+    ("h", "w", "hg", "wg"),
+    [(80, 80, 300, 3), (40, 40, 300, 6), (20, 20, 300, 3), (13, 17, 50, 4)],
+)
+def test_bilinear_grid_sample__matches_grid_sample(
+    h: int, w: int, hg: int, wg: int
+) -> None:
+    # The gather-based implementation must match F.grid_sample (bilinear, zero
+    # padding, align_corners=False), including for out-of-bounds coordinates which
+    # deformable attention produces.
+    torch.manual_seed(0)
+    im = torch.randn(4, 12, h, w)
+    grid = torch.rand(4, hg, wg, 2) * 2.6 - 1.3  # spans outside [-1, 1]
+
+    expected = F.grid_sample(
+        im, grid, mode="bilinear", padding_mode="zeros", align_corners=False
+    )
+    got = bilinear_grid_sample(im, grid, align_corners=False)
+
+    torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_deformable_attention_core__bilinear_gather_matches_default() -> None:
+    # The "bilinear_gather" method (used for TensorRT-safe export) must reproduce
+    # the default grid_sample-based sampling numerically.
+    torch.manual_seed(0)
+    bs, n_head, c = 2, 8, 16
+    value_spatial_shapes = [(20, 20), (10, 10)]
+    num_points_list = [4, 4]
+    len_q = 30
+    value_len = sum(h * w for h, w in value_spatial_shapes)
+
+    value = torch.randn(bs, value_len, n_head, c)
+    sampling_locations = torch.rand(bs, len_q, n_head, sum(num_points_list), 2)
+    attention_weights = torch.rand(bs, len_q, n_head, sum(num_points_list))
+
+    default = deformable_attention_core_func_v2(
+        method="default",
+        value=value,
+        value_spatial_shapes=value_spatial_shapes,
+        sampling_locations=sampling_locations,
+        attention_weights=attention_weights,
+        num_points_list=num_points_list,
+    )
+    gather = deformable_attention_core_func_v2(
+        method="bilinear_gather",
+        value=value,
+        value_spatial_shapes=value_spatial_shapes,
+        sampling_locations=sampling_locations,
+        attention_weights=attention_weights,
+        num_points_list=num_points_list,
+    )
+
+    torch.testing.assert_close(gather, default, atol=1e-5, rtol=1e-5)