Build spatial_shapes from symbolic Shape ops for TensorRT-safe ONNX export (roboflow#1155)

- Fixed TensorRT-safe ONNX export for Transformer `spatial_shapes`, avoiding ScatterND in shape-tensor computation.
- Updated ONNX export behavior to preserve dynamic-batch compatibility across detection, segmentation, and keypoint models.
- Improved Transformer export validation with coverage for dynamic batches, single-level feature maps, non-square spatial shapes, and PyTorch API compatibility.

---------

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Author: 5 people

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - `HungarianMatcher.forward()` now uses the configured `focal_alpha` in the focal classification matching cost. Previously the value was hardcoded to `0.25`, silently ignoring any non-default `focal_alpha` passed to the constructor or `build_matcher`. This misaligned the bipartite matching cost with the focal classification loss in `criterion.py`, which correctly used `self.focal_alpha`. ([#1147](https://github.com/roboflow/rf-detr/pull/1147))
+- `spatial_shapes` in `Transformer.forward()` is now built from symbolic `Shape` ops (`torch.stack` of per-level `torch._shape_as_tensor` slices) instead of `torch.empty` + in-place index assignment. The previous pattern emitted a `ScatterND` feeding a shape tensor (`level_start_index`), which TensorRT rejected with "IScatterLayer cannot be used to compute a shape tensor". This fix is required to export any RF-DETR model to a TensorRT engine. ([#1155](https://github.com/roboflow/rf-detr/pull/1155))
 
 ---
 

src/rfdetr/models/transformer.py

@@ -284,16 +284,10 @@ def forward(
         src_flatten = []
         mask_flatten = [] if masks is not None else None
         lvl_pos_embed_flatten = []
-        # Build spatial_shapes as a tensor directly so that the ONNX tracer
-        # can track h/w symbolically instead of baking constants into a constant
-        # shape node.
-        spatial_shapes = torch.empty((len(srcs), 2), device=srcs[0].device, dtype=torch.long)
         spatial_shapes_hw: list[tuple[int, int]] = []
         valid_ratios = [] if masks is not None else None
         for lvl, (src, pos_embed) in enumerate(zip(srcs, pos_embeds)):
             _, c, h, w = src.shape
-            spatial_shapes[lvl, 0] = h
-            spatial_shapes[lvl, 1] = w
             spatial_shapes_hw.append((h, w))
 
             src = src.flatten(2).transpose(1, 2)  # bs, hw, c
@@ -309,6 +303,18 @@ def forward(
             mask_flatten = torch.cat(mask_flatten, 1)  # bs, \sum{hxw}
             valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
         lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)  # bs, \sum{hxw}, c
+        # spatial_shapes must not be built by torch.empty(...) + in-place index assignment:
+        # that emits a ScatterND feeding a shape tensor (level_start_index), which TensorRT
+        # rejects ("IScatterLayer cannot be used to compute a shape tensor").
+        # torch.as_tensor(python-int list) avoids ScatterND but bakes values as a Constant.
+        # torch.stack of per-level torch._shape_as_tensor slices also produces a Constant
+        # node in TorchScript ONNX export (the tracer records concrete H,W values at trace
+        # time), but that Constant is accepted by TensorRT as a valid shape tensor source —
+        # unlike ScatterND. torch._shape_as_tensor(t) is a private ATen op that returns a
+        # 1-D int64 tensor of t's dimension sizes; [2:4] extracts (H, W) from NCHW.
+        spatial_shapes = torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs]).to(
+            device=srcs[0].device, dtype=torch.long
+        )
         level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
 
         # Flatten optional dual-projector features for keypoint-specific cross-attention.