diff --git a/CHANGELOG.md b/CHANGELOG.md index 91632ece7..c3ee433d1 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/lightly_train/_export/export_onnx.py b/src/lightly_train/_export/export_onnx.py index 19aeb9fc0..743f25937 100644 --- a/src/lightly_train/_export/export_onnx.py +++ b/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, diff --git a/src/lightly_train/_export/onnx_helpers.py b/src/lightly_train/_export/onnx_helpers.py index ba2b92789..5910bb6f6 100644 --- a/src/lightly_train/_export/onnx_helpers.py +++ b/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. diff --git a/src/lightly_train/_export/tensorrt_helpers.py b/src/lightly_train/_export/tensorrt_helpers.py index 5ae60df09..3531cd041 100644 --- a/src/lightly_train/_export/tensorrt_helpers.py +++ b/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) diff --git a/src/lightly_train/_task_models/ltdetr_instance_segmentation/task_model.py b/src/lightly_train/_task_models/ltdetr_instance_segmentation/task_model.py index 8fb0e37a0..01ce49eb3 100644 --- a/src/lightly_train/_task_models/ltdetr_instance_segmentation/task_model.py +++ b/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 diff --git a/src/lightly_train/_task_models/object_detection_components/dfine_decoder.py b/src/lightly_train/_task_models/object_detection_components/dfine_decoder.py index 9d2a8f337..5ed4a125b 100644 --- a/src/lightly_train/_task_models/object_detection_components/dfine_decoder.py +++ b/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, diff --git a/src/lightly_train/_task_models/object_detection_components/rtdetrv2_decoder.py b/src/lightly_train/_task_models/object_detection_components/rtdetrv2_decoder.py index fe956ea7b..79ba0545f 100644 --- a/src/lightly_train/_task_models/object_detection_components/rtdetrv2_decoder.py +++ b/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, diff --git a/src/lightly_train/_task_models/object_detection_components/utils.py b/src/lightly_train/_task_models/object_detection_components/utils.py index 94175b9cd..358f01e1b 100644 --- a/src/lightly_train/_task_models/object_detection_components/utils.py +++ b/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 = ( diff --git a/tests/_task_models/ltdetr_instance_segmentation/test_task_model.py b/tests/_task_models/ltdetr_instance_segmentation/test_task_model.py index 63726913c..5317bf371 100644 --- a/tests/_task_models/ltdetr_instance_segmentation/test_task_model.py +++ b/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) diff --git a/tests/_task_models/ltdetr_object_detection/test_task_model.py b/tests/_task_models/ltdetr_object_detection/test_task_model.py index 620360d0a..1a6585b12 100644 --- a/tests/_task_models/ltdetr_object_detection/test_task_model.py +++ b/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( diff --git a/tests/_task_models/object_detection_components/test_utils.py b/tests/_task_models/object_detection_components/test_utils.py index 54156c9e8..0b158c9b9 100644 --- a/tests/_task_models/object_detection_components/test_utils.py +++ b/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)