diff --git a/apps/benchmark/performance/release.yaml b/apps/benchmark/performance/release.yaml index 8f71ee2901..1ebb281f96 100644 --- a/apps/benchmark/performance/release.yaml +++ b/apps/benchmark/performance/release.yaml @@ -54,6 +54,10 @@ excluded_profiles: performance baseline runs Ultralytics, which cannot load a legacy YOLOv5 archive: those archives pickle classes from the standalone yolov5 repository, which is not a dependency here. + - model: yolox-s + reason: >- + Functional and official-reference parity are covered by family tests, + but no matching release-performance workload or receipt is provided. - model: mobilenetv4-conv-small reason: &mobilenetv4_performance_exclusion >- Functional and timm reference-parity qualification is present for every diff --git a/families/yolox/__init__.py b/families/yolox/__init__.py new file mode 100644 index 0000000000..2d5b3dd219 --- /dev/null +++ b/families/yolox/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""YOLOX object detection family.""" diff --git a/families/yolox/checkpoint.py b/families/yolox/checkpoint.py new file mode 100644 index 0000000000..97b3c3eac6 --- /dev/null +++ b/families/yolox/checkpoint.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read an official YOLOX state dictionary without unpickling model code.""" + +from pathlib import Path + +import numpy as np +import torch + +from .support import ARCHIVES + + +class Checkpoint: + def __init__(self, state: dict[str, torch.Tensor], *, image_size: int = 640) -> None: + if not isinstance(state, dict) or not state: + raise ValueError("YOLOX checkpoint must contain a non-empty model state dictionary") + self.state = state + self.image_size = image_size + self.used: set[str] = set() + for name, tensor in state.items(): + if not isinstance(name, str) or not isinstance(tensor, torch.Tensor): + raise ValueError("YOLOX model state must contain named tensors") + if not torch.isfinite(tensor).all(): + raise ValueError(f"YOLOX checkpoint contains non-finite values: {name}") + + @classmethod + def open(cls, model_dir: Path) -> "Checkpoint": + paths = [model_dir / name for name in ARCHIVES if (model_dir / name).is_file()] + if len(paths) != 1: + raise ValueError("YOLOX model directory must contain exactly one official checkpoint") + path = paths[0] + archive = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(archive, dict) or "model" not in archive: + raise ValueError("YOLOX checkpoint must contain a model state dictionary") + # Training/evaluation resolution is not stored in a state dictionary. + image_size = 416 if path.name in {"yolox_nano.pth", "yolox_tiny.pth"} else 640 + return cls(archive["model"], image_size=image_size) + + def tensor(self, name: str) -> np.ndarray: + if name not in self.state: + raise ValueError(f"YOLOX checkpoint is missing {name}") + self.used.add(name) + return self.state[name].detach().float().numpy() + + def assert_consumed(self) -> None: + unused = set(self.state) - self.used + unused = {name for name in unused if not name.endswith(".bn.num_batches_tracked")} + if unused: + raise ValueError(f"Unsupported YOLOX checkpoint tensors: {sorted(unused)}") diff --git a/families/yolox/graph.py b/families/yolox/graph.py new file mode 100644 index 0000000000..3de4418e18 --- /dev/null +++ b/families/yolox/graph.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small TensorRT graph vocabulary owned by YOLOX.""" + +from __future__ import annotations + +import numpy as np +import tensorrt as trt + + +def convolution( + network, + tensor, + weight: np.ndarray, + bias: np.ndarray, + *, + stride: int = 1, + padding: int = 0, + groups: int = 1, + dtype: np.dtype, +): + layer = network.add_convolution_nd( + tensor, + num_output_maps=int(weight.shape[0]), + kernel_shape=(int(weight.shape[2]), int(weight.shape[3])), + kernel=trt.Weights(np.ascontiguousarray(weight, dtype=dtype)), + bias=trt.Weights(np.ascontiguousarray(bias, dtype=dtype)), + ) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX convolution") + layer.stride_nd = (stride, stride) + layer.padding_nd = (padding, padding) + layer.num_groups = groups + return layer.get_output(0) + + +def silu(network, tensor): + """SiLU with FP32 internal arithmetic, as in PyTorch's half kernel. + + Rounding sigmoid and its product separately in FP16 differs from the + upstream single activation and accumulates across the CSP blocks. + """ + dtype = tensor.dtype + if dtype == trt.float16: + cast = network.add_cast(tensor, trt.float32) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX SiLU input cast") + tensor = cast.get_output(0) + gate = network.add_activation(tensor, trt.ActivationType.SIGMOID) + if gate is None: + raise RuntimeError("TensorRT rejected a YOLOX SiLU sigmoid") + product = network.add_elementwise(tensor, gate.get_output(0), trt.ElementWiseOperation.PROD) + if product is None: + raise RuntimeError("TensorRT rejected a YOLOX SiLU product") + output = product.get_output(0) + if dtype == trt.float16: + cast = network.add_cast(output, trt.float16) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX SiLU output cast") + output = cast.get_output(0) + return output + + +def leaky_relu(network, tensor): + layer = network.add_activation(tensor, trt.ActivationType.LEAKY_RELU) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX leaky ReLU") + layer.alpha = 0.1 + return layer.get_output(0) + + +def add(network, left, right): + layer = network.add_elementwise(left, right, trt.ElementWiseOperation.SUM) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX add") + return layer.get_output(0) + + +def concatenate(network, tensors, *, axis: int = 1): + layer = network.add_concatenation(tensors) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX concatenation") + layer.axis = axis + return layer.get_output(0) + + +def slice_axis(network, tensor, *, axis: int, start: int, count: int): + shape = [int(value) for value in tensor.shape] + starts, sizes = [0] * len(shape), list(shape) + starts[axis], sizes[axis] = start, count + layer = network.add_slice(tensor, trt.Dims(starts), trt.Dims(sizes), trt.Dims([1] * len(shape))) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX slice") + return layer.get_output(0) + + +def max_pool(network, tensor, *, kernel: int, stride: int, padding: int): + layer = network.add_pooling_nd(tensor, trt.PoolingType.MAX, (kernel, kernel)) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX max pool") + layer.stride_nd = (stride, stride) + layer.padding_nd = (padding, padding) + return layer.get_output(0) + + +def nearest_upsample(network, tensor, factor: int): + layer = network.add_resize(tensor) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX upsample") + layer.resize_mode = trt.InterpolationMode.NEAREST + layer.scales = [1.0, 1.0, float(factor), float(factor)] + return layer.get_output(0) + + +def reshape(network, tensor, shape: tuple[int, ...]): + layer = network.add_shuffle(tensor) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX reshape") + layer.reshape_dims = trt.Dims(shape) + return layer.get_output(0) + + +def permute(network, tensor, permutation: tuple[int, ...]): + layer = network.add_shuffle(tensor) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX permutation") + layer.second_transpose = trt.Permutation(permutation) + return layer.get_output(0) + + +def sigmoid(network, tensor): + layer = network.add_activation(tensor, trt.ActivationType.SIGMOID) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX sigmoid") + return layer.get_output(0) + + +def scale(network, tensor, factor: float, *, dtype: np.dtype): + shape = (1,) * len(tuple(tensor.shape)) + layer = network.add_constant(shape, trt.Weights(np.array([factor], dtype=dtype).reshape(shape))) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX scale constant") + values = layer.get_output(0) + if values.dtype != tensor.dtype: + cast = network.add_cast(values, tensor.dtype) + if cast is None: + raise RuntimeError("TensorRT rejected a YOLOX scale cast") + values = cast.get_output(0) + product = network.add_elementwise(tensor, values, trt.ElementWiseOperation.PROD) + if product is None: + raise RuntimeError("TensorRT rejected a YOLOX scale product") + return product.get_output(0) + + +def constant(network, values: np.ndarray, *, dtype: np.dtype, like=None): + layer = network.add_constant( + values.shape, trt.Weights(np.ascontiguousarray(values, dtype=dtype)) + ) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX constant") + output = layer.get_output(0) + if like is None or output.dtype == like.dtype: + return output + cast = network.add_cast(output, like.dtype) + if cast is None: + raise RuntimeError("TensorRT rejected a YOLOX constant cast") + return cast.get_output(0) + + +def subtract(network, left, right): + layer = network.add_elementwise(left, right, trt.ElementWiseOperation.SUB) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX subtraction") + return layer.get_output(0) + + +def top_k(network, tensor, *, k: int, axis: int): + """Largest `k` values along one axis, with their indices.""" + layer = network.add_topk(tensor, trt.TopKOperation.MAX, k, 1 << axis) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX top-k") + return layer.get_output(0), layer.get_output(1) + + +def multiply(network, left, right): + """Element-wise product; TensorRT broadcasts size-one axes.""" + layer = network.add_elementwise(left, right, trt.ElementWiseOperation.PROD) + if layer is None: + raise RuntimeError("TensorRT rejected a YOLOX product") + return layer.get_output(0) diff --git a/families/yolox/model.py b/families/yolox/model.py new file mode 100644 index 0000000000..267ff497d3 --- /dev/null +++ b/families/yolox/model.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""YOLOX: CSPDarknet/PAFPN or Darknet/FPN and a decoupled anchor-free head. + +The topology follows Megvii-BaseDetection/YOLOX at +6ddff4824372906469a7fae2dc3206c7aa4bbaee, exps/default/. +TensorRT owns lowering and execution; this family specifies the graph. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import tensorrt as trt + +from . import graph +from .checkpoint import Checkpoint + +if TYPE_CHECKING: + from tensorrt_model_connect.build import BuildRequest + from tensorrt_model_connect.bundle_writer import BundleWriter + + +# Exp.get_model overrides the PyTorch default epsilon before loading weights. +_BATCH_NORM_EPSILON = 1e-3 +_NUM_CLASSES = 80 +_STRIDES = (8, 16, 32) + + +def _fold(checkpoint: Checkpoint, prefix: str, dtype: np.dtype) -> tuple[np.ndarray, np.ndarray]: + weight = checkpoint.tensor(f"{prefix}.conv.weight") + gamma = checkpoint.tensor(f"{prefix}.bn.weight") + beta = checkpoint.tensor(f"{prefix}.bn.bias") + mean = checkpoint.tensor(f"{prefix}.bn.running_mean") + variance = checkpoint.tensor(f"{prefix}.bn.running_var") + if weight.ndim != 4 or any( + v.shape != (weight.shape[0],) for v in (gamma, beta, mean, variance) + ): + raise ValueError(f"YOLOX convolution/BatchNorm shape mismatch: {prefix}") + if np.any(variance < 0): + raise ValueError(f"YOLOX BatchNorm variance must be non-negative: {prefix}") + scale = gamma / np.sqrt(variance + _BATCH_NORM_EPSILON) + return (weight * scale.reshape(-1, 1, 1, 1)).astype(dtype), (beta - mean * scale).astype(dtype) + + +class _Weights: + """Folded convolutions and plain tensors, addressed by checkpoint prefix.""" + + def __init__(self, checkpoint: Checkpoint, dtype: np.dtype) -> None: + self._checkpoint = checkpoint + self._dtype = dtype + self._folded: dict[str, tuple[np.ndarray, np.ndarray]] = {} + + def conv(self, prefix: str) -> tuple[np.ndarray, np.ndarray]: + if prefix not in self._folded: + self._folded[prefix] = _fold(self._checkpoint, prefix, self._dtype) + return self._folded[prefix] + + def raw(self, name: str) -> np.ndarray: + return self._checkpoint.tensor(name).astype(self._dtype) + + def exists(self, name: str) -> bool: + return name in self._checkpoint.state + + +def _conv(network, tensor, weights: _Weights, prefix: str, dtype, *, stride: int = 1): + if weights.exists(f"{prefix}.dconv.conv.weight"): + # Rounding between depthwise and pointwise convolutions amplifies score error. + if dtype == np.float16: + cast = network.add_cast(tensor, trt.float32) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX depthwise input cast") + tensor = cast.get_output(0) + tensor = _conv(network, tensor, weights, f"{prefix}.dconv", np.float32, stride=stride) + tensor = _conv(network, tensor, weights, f"{prefix}.pconv", np.float32) + if dtype == np.float16: + cast = network.add_cast(tensor, trt.float16) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX depthwise output cast") + tensor = cast.get_output(0) + return tensor + weight, bias = weights.conv(prefix) + groups = int(tensor.shape[1]) if prefix.endswith(".dconv") else 1 + if weight.shape[1] * groups != int(tensor.shape[1]): + raise ValueError(f"YOLOX input channel mismatch: {prefix}") + tensor = graph.convolution( + network, + tensor, + weight, + bias, + stride=stride, + padding=weight.shape[2] // 2, + groups=groups, + dtype=dtype, + ) + if weights.exists("backbone.backbone.stem.0.conv.weight"): + return graph.leaky_relu(network, tensor) + return graph.silu(network, tensor) + + +def _csp(network, tensor, weights: _Weights, prefix: str, dtype, *, residual: bool): + left = _conv(network, tensor, weights, f"{prefix}.conv1", dtype) + right = _conv(network, tensor, weights, f"{prefix}.conv2", dtype) + # Preserve the residual path around depthwise bottlenecks in FP32 as well. + inner_dtype = np.float32 if weights.exists(f"{prefix}.m.0.conv2.dconv.conv.weight") else dtype + if inner_dtype != dtype: + cast = network.add_cast(left, trt.float32) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX bottleneck input cast") + left = cast.get_output(0) + index = 0 + while weights.exists(f"{prefix}.m.{index}.conv1.conv.weight"): + inner = _conv(network, left, weights, f"{prefix}.m.{index}.conv1", inner_dtype) + inner = _conv(network, inner, weights, f"{prefix}.m.{index}.conv2", inner_dtype) + left = graph.add(network, left, inner) if residual else inner + index += 1 + if index == 0: + raise ValueError(f"YOLOX CSP block has no bottlenecks: {prefix}") + if inner_dtype != dtype: + cast = network.add_cast(left, right.dtype) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX bottleneck output cast") + left = cast.get_output(0) + return _conv( + network, graph.concatenate(network, [left, right]), weights, f"{prefix}.conv3", dtype + ) + + +def _focus(network, tensor): + # Order is top-left, bottom-left, top-right, bottom-right, not row-major. + batch, channels, height, width = map(int, tensor.shape) + parts = [] + for y, x in ((0, 0), (1, 0), (0, 1), (1, 1)): + layer = network.add_slice( + tensor, (0, 0, y, x), (batch, channels, height // 2, width // 2), (1, 1, 2, 2) + ) + if layer is None: + raise RuntimeError("TensorRT rejected YOLOX Focus") + parts.append(layer.get_output(0)) + return graph.concatenate(network, parts) + + +def _spp(network, tensor, weights: _Weights, prefix: str, dtype): + entry = _conv(network, tensor, weights, f"{prefix}.conv1", dtype) + # Parallel SPP pools in the order used by the official implementation. + parts = [entry] + [ + graph.max_pool(network, entry, kernel=k, stride=1, padding=k // 2) for k in (5, 9, 13) + ] + return _conv(network, graph.concatenate(network, parts), weights, f"{prefix}.conv2", dtype) + + +def _darknet(network, pixels, weights: _Weights, dtype): + prefix = "backbone.backbone" + tensor = _conv(network, pixels, weights, f"{prefix}.stem.0", np.float32) + if dtype == np.float16: + cast = network.add_cast(tensor, trt.float16) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX Darknet feature cast") + tensor = cast.get_output(0) + outputs = [] + for stage in ("stem", "dark2", "dark3", "dark4", "dark5"): + index = 1 if stage == "stem" else 0 + block = f"{prefix}.{stage}" + tensor = _conv(network, tensor, weights, f"{block}.{index}", dtype, stride=2) + index += 1 + while weights.exists(f"{block}.{index}.layer1.conv.weight"): + inner = _conv(network, tensor, weights, f"{block}.{index}.layer1", dtype) + inner = _conv(network, inner, weights, f"{block}.{index}.layer2", dtype) + tensor = graph.add(network, tensor, inner) + index += 1 + if stage == "dark5": + tensor = _conv(network, tensor, weights, f"{block}.{index}", dtype) + tensor = _conv(network, tensor, weights, f"{block}.{index + 1}", dtype) + tensor = _spp(network, tensor, weights, f"{block}.{index + 2}", dtype) + tensor = _conv(network, tensor, weights, f"{block}.{index + 3}", dtype) + tensor = _conv(network, tensor, weights, f"{block}.{index + 4}", dtype) + if stage in {"dark3", "dark4", "dark5"}: + outputs.append(tensor) + return outputs + + +def _backbone(network, pixels, weights: _Weights, dtype): + if weights.exists("backbone.backbone.stem.0.conv.weight"): + return _darknet(network, pixels, weights, dtype) + prefix = "backbone.backbone" + # Preserve small color differences in the unnormalized BGR byte input. + tensor = _conv(network, _focus(network, pixels), weights, f"{prefix}.stem.conv", np.float32) + if dtype == np.float16: + cast = network.add_cast(tensor, trt.float16) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX feature cast") + tensor = cast.get_output(0) + outputs = [] + for stage in range(2, 6): + tensor = _conv(network, tensor, weights, f"{prefix}.dark{stage}.0", dtype, stride=2) + if stage == 5: + tensor = _spp(network, tensor, weights, f"{prefix}.dark5.1", dtype) + tensor = _csp( + network, + tensor, + weights, + f"{prefix}.dark{stage}.{2 if stage == 5 else 1}", + dtype, + residual=stage != 5, + ) + if stage >= 3: + outputs.append(tensor) + return outputs + + +def _fpn(network, sources, weights: _Weights, dtype): + dark3, dark4, tensor = sources + outputs = [tensor] + for level, source in enumerate((dark4, dark3), start=1): + tensor = _conv(network, tensor, weights, f"backbone.out{level}_cbl", dtype) + tensor = graph.concatenate(network, [graph.nearest_upsample(network, tensor, 2), source]) + index = 0 + while weights.exists(f"backbone.out{level}.{index}.conv.weight"): + tensor = _conv(network, tensor, weights, f"backbone.out{level}.{index}", dtype) + index += 1 + outputs.append(tensor) + return tuple(reversed(outputs)) + + +def _neck(network, sources, weights: _Weights, dtype): + # PAFPN and the head need FP32 to keep score error within 0.01. + promoted = [] + for source in sources: + if source.dtype != trt.float32: + cast = network.add_cast(source, trt.float32) + if cast is None: + raise RuntimeError("TensorRT rejected the YOLOX PAFPN feature cast") + source = cast.get_output(0) + promoted.append(source) + if weights.exists("backbone.out1_cbl.conv.weight"): + return _fpn(network, promoted, weights, dtype) + dark3, dark4, dark5 = promoted + lateral = _conv(network, dark5, weights, "backbone.lateral_conv0", dtype) + merged = graph.concatenate(network, [graph.nearest_upsample(network, lateral, 2), dark4]) + upper = _csp(network, merged, weights, "backbone.C3_p4", dtype, residual=False) + reduced = _conv(network, upper, weights, "backbone.reduce_conv1", dtype) + merged = graph.concatenate(network, [graph.nearest_upsample(network, reduced, 2), dark3]) + p3 = _csp(network, merged, weights, "backbone.C3_p3", dtype, residual=False) + merged = graph.concatenate( + network, [_conv(network, p3, weights, "backbone.bu_conv2", dtype, stride=2), reduced] + ) + p4 = _csp(network, merged, weights, "backbone.C3_n3", dtype, residual=False) + merged = graph.concatenate( + network, [_conv(network, p4, weights, "backbone.bu_conv1", dtype, stride=2), lateral] + ) + p5 = _csp(network, merged, weights, "backbone.C3_n4", dtype, residual=False) + return p3, p4, p5 + + +def _predict(network, tensor, weights: _Weights, prefix: str, dtype, *, channels: int): + weight, bias = weights.raw(f"{prefix}.weight"), weights.raw(f"{prefix}.bias") + if weight.shape != (channels, int(tensor.shape[1]), 1, 1) or bias.shape != (channels,): + raise ValueError(f"Unsupported YOLOX prediction shape: {prefix}") + return graph.convolution(network, tensor, weight, bias, dtype=dtype) + + +def _detect(network, sources, weights: _Weights, dtype): + box_parts, score_parts = [], [] + for level, (tensor, stride) in enumerate(zip(sources, _STRIDES, strict=True)): + stem = _conv(network, tensor, weights, f"head.stems.{level}", dtype) + cls_feature, reg_feature = stem, stem + for index in range(2): + cls_feature = _conv( + network, cls_feature, weights, f"head.cls_convs.{level}.{index}", dtype + ) + reg_feature = _conv( + network, reg_feature, weights, f"head.reg_convs.{level}.{index}", dtype + ) + regression = _predict( + network, reg_feature, weights, f"head.reg_preds.{level}", dtype, channels=4 + ) + objectness = _predict( + network, reg_feature, weights, f"head.obj_preds.{level}", dtype, channels=1 + ) + classes = _predict( + network, cls_feature, weights, f"head.cls_preds.{level}", dtype, channels=_NUM_CLASSES + ) + rows, columns = map(int, regression.shape[2:]) + cells = rows * columns + regression = graph.reshape(network, regression, (4, cells)) + offset = graph.slice_axis(network, regression, axis=0, start=0, count=2) + log_size = graph.slice_axis(network, regression, axis=0, start=2, count=2) + y, x = np.meshgrid(np.arange(rows), np.arange(columns), indexing="ij") + grid = graph.constant(network, np.stack([x.ravel(), y.ravel()]), dtype=np.float32) + centre = graph.scale(network, graph.add(network, offset, grid), stride, dtype=np.float32) + exp = network.add_unary(log_size, trt.UnaryOperation.EXP) + if exp is None: + raise RuntimeError("TensorRT rejected YOLOX box exponential") + half = graph.scale(network, exp.get_output(0), stride * 0.5, dtype=np.float32) + corners = graph.concatenate( + network, + [graph.subtract(network, centre, half), graph.add(network, centre, half)], + axis=0, + ) + box_parts.append(graph.permute(network, corners, (1, 0))) + probabilities = graph.reshape( + network, graph.sigmoid(network, classes), (_NUM_CLASSES, cells) + ) + confidence = graph.reshape(network, graph.sigmoid(network, objectness), (1, cells)) + score_parts.append( + graph.permute(network, graph.multiply(network, probabilities, confidence), (1, 0)) + ) + boxes = graph.concatenate(network, box_parts, axis=0) + scores = graph.concatenate(network, score_parts, axis=0) + best, index = graph.top_k(network, scores, k=1, axis=1) + total = int(boxes.shape[0]) + return (boxes, graph.reshape(network, best, (total,)), graph.reshape(network, index, (total,))) + + +def _build_engine(checkpoint: Checkpoint, precision: str, verbose: bool) -> bytes: + if precision not in {"fp32", "fp16"}: + raise ValueError(f"Unsupported YOLOX precision: {precision}") + numpy_dtype = np.float16 if precision == "fp16" else np.float32 + # Fold once in FP32; graph.convolution converts weights to each layer's dtype. + weights = _Weights(checkpoint, np.float32) + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + config.clear_flag(trt.BuilderFlag.TF32) + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30) + size = checkpoint.image_size + pixels = network.add_input("pixel_values", trt.float32, (1, 3, size, size)) + if pixels is None: + raise RuntimeError("TensorRT rejected the YOLOX input") + sources = _backbone(network, pixels, weights, numpy_dtype) + sources = _neck(network, sources, weights, np.float32) + outputs = _detect(network, sources, weights, np.float32) + checkpoint.assert_consumed() + for tensor, name in zip(outputs, ("boxes", "scores", "classes"), strict=True): + tensor.name = name + network.mark_output(tensor) + plan = builder.build_serialized_network(network, config) + if plan is None: + raise RuntimeError("TensorRT YOLOX engine build failed") + return bytes(plan) + + +def build(request: "BuildRequest", writer: "BundleWriter") -> None: + """Build an official 80-class YOLOX detector at its published input size.""" + if request.backend != "trt": + raise NotImplementedError("yolox supports only backend=trt") + if request.task != "object_detection": + raise ValueError("yolox supports only task=object_detection") + if request.dynamic_kv_cache: + raise NotImplementedError("yolox does not support dynamic_kv_cache") + if request.image_height is not None or request.image_width is not None: + raise NotImplementedError("yolox does not support image_height or image_width overrides") + if request.video_num_frames is not None: + raise NotImplementedError("yolox does not support video_num_frames") + if request.max_batch_size != 1: + raise NotImplementedError("yolox does not support max_batch_size other than one") + if request.tensor_parallel_size != 1: + raise NotImplementedError("yolox does not support tensor parallelism") + if request.context_parallel_size != 1: + raise NotImplementedError("yolox does not support context parallelism") + if request.quantization not in {None, "none"}: + raise NotImplementedError("yolox does not support quantization") + if request.fp32_layers: + raise NotImplementedError("yolox does not support mixed-precision layer overrides") + if request.max_sequence_length not in {None, 1}: + raise NotImplementedError("yolox does not support max_sequence_length") + checkpoint = Checkpoint.open(Path(request.model_dir)) + plan = _build_engine(checkpoint, str(request.precision).lower(), bool(request.verbose)) + writer.set_header(family="yolox", task=request.task, backend=request.backend) + writer.add_bytes("engine.plan", plan) + writer.add_json( + "runtime.json", + { + "input_image_h": checkpoint.image_size, + "input_image_w": checkpoint.image_size, + "pad_value": 114, + "score_threshold": 0.25, + "iou_threshold": 0.45, + "num_classes": _NUM_CLASSES, + "max_detections": sum((checkpoint.image_size // stride) ** 2 for stride in _STRIDES), + }, + ) diff --git a/families/yolox/requirements.txt b/families/yolox/requirements.txt new file mode 100644 index 0000000000..5d1d0caf68 --- /dev/null +++ b/families/yolox/requirements.txt @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# PyTorch is used to read the checkpoint; the remaining packages run the official reference. +torch>=2.6 +torchvision +opencv-python-headless +loguru +psutil +pycocotools +Pillow +tabulate +tqdm diff --git a/families/yolox/runtime/CMakeLists.txt b/families/yolox/runtime/CMakeLists.txt new file mode 100644 index 0000000000..d50ac0f99c --- /dev/null +++ b/families/yolox/runtime/CMakeLists.txt @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +add_library(trtmc_model_yolox SHARED + image_preprocess_seam.cpp + pipeline.cpp + plugin.cpp +) +target_include_directories(trtmc_model_yolox PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/core/runtime/include +) +target_include_directories(trtmc_model_yolox SYSTEM PRIVATE + ${TRTMC_CUDA_INCLUDE_DIR} +) +target_link_libraries(trtmc_model_yolox PRIVATE + trtmc_core + nlohmann_json::nlohmann_json + ${TRTMC_CUDART_LIBRARY} +) +target_compile_options(trtmc_model_yolox PRIVATE + "$<$:-Wall;-Wextra;-Wpedantic>" +) +set_target_properties(trtmc_model_yolox PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" + BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" +) +install(TARGETS trtmc_model_yolox + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} +) + +if(TRTMC_BUILD_TESTS) + add_executable(test_yolox_image_preprocess + ${PROJECT_SOURCE_DIR}/families/yolox/tests/cpp/test_image_preprocess_seam.cpp + ) + target_include_directories(test_yolox_image_preprocess PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_link_libraries(test_yolox_image_preprocess PRIVATE + trtmc_model_yolox + trtmc_core + ${TRTMC_CUDART_LIBRARY} + ) + target_compile_options(test_yolox_image_preprocess PRIVATE + -Wall -Wextra -Wpedantic + ) + add_test( + NAME yolox_image_preprocess + COMMAND test_yolox_image_preprocess + ) +endif() diff --git a/families/yolox/runtime/image_preprocess_seam.cpp b/families/yolox/runtime/image_preprocess_seam.cpp new file mode 100644 index 0000000000..b4a81161d5 --- /dev/null +++ b/families/yolox/runtime/image_preprocess_seam.cpp @@ -0,0 +1,77 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/yolox/runtime/image_preprocess_seam.h" + +#include +#include +#include + +namespace trtmc { +namespace { + +float sample_bilinear(const float* image, std::int32_t height, std::int32_t width, + std::int32_t channel, float y, float x) { + // The task supplies interleaved RGB floats. YOLOX resizes integer bytes, + // so quantize the four input samples before interpolation. + const float clamped_y = std::clamp(y, 0.0F, static_cast(height - 1)); + const float clamped_x = std::clamp(x, 0.0F, static_cast(width - 1)); + const auto y0 = static_cast(clamped_y); + const auto x0 = static_cast(clamped_x); + const std::int32_t y1 = std::min(y0 + 1, height - 1); + const std::int32_t x1 = std::min(x0 + 1, width - 1); + const double wy = static_cast(clamped_y) - y0; + const double wx = static_cast(clamped_x) - x0; + const auto at = [&](std::int32_t row, std::int32_t column) { + const auto index = (static_cast(row) * width + column) * 3U + channel; + const float value = image[index]; + if (!std::isfinite(value) || value < 0.0F || value > 1.0F) + throw std::invalid_argument("YOLOX input must be finite RGB pixels in [0, 1]"); + return std::round(value * 255.0F); + }; + const double top = at(y0, x0) * (1.0 - wx) + at(y0, x1) * wx; + const double bottom = at(y1, x0) * (1.0 - wx) + at(y1, x1) * wx; + // OpenCV's fixed-point byte resize may differ by one; E2E bounds this. + return static_cast(std::round(top * (1.0 - wy) + bottom * wy)); +} + +} // namespace + +std::vector preprocess_yolox_image(const float* pixels, std::int32_t height, + std::int32_t width, const YoloxPreprocessConfig& config, + YoloxLetterbox& letterbox) { + if (pixels == nullptr || height <= 0 || width <= 0) + throw std::invalid_argument("YOLOX preprocessing needs a non-empty image"); + if (config.input_image_h <= 0 || config.input_image_w <= 0 || + !std::isfinite(config.pad_value) || config.pad_value < 0.0F || config.pad_value > 255.0F) + throw std::invalid_argument("YOLOX preprocessing configuration is invalid"); + + // Upstream non-legacy preproc: truncate resized dimensions, paste at the + // top left, pad the right and bottom with 114, keep BGR bytes without /255. + const double scale = std::min(static_cast(config.input_image_h) / height, + static_cast(config.input_image_w) / width); + const auto scaled_h = static_cast(height * scale); + const auto scaled_w = static_cast(width * scale); + if (scaled_h < 1 || scaled_w < 1) + throw std::invalid_argument("YOLOX image aspect ratio produces an empty resize"); + letterbox = {static_cast(scale), 0.0F, 0.0F}; + const auto plane = static_cast(config.input_image_h) * config.input_image_w; + std::vector values(3U * plane, config.pad_value); + for (std::int32_t channel = 0; channel < 3; ++channel) { + float* target = values.data() + static_cast(channel) * plane; + for (std::int32_t row = 0; row < scaled_h; ++row) { + const float y = static_cast((row + 0.5) * height / scaled_h - 0.5); + for (std::int32_t column = 0; column < scaled_w; ++column) { + const float x = static_cast((column + 0.5) * width / scaled_w - 0.5); + // Select the RGB source channel for this BGR output plane. + target[static_cast(row) * config.input_image_w + column] = + sample_bilinear(pixels, height, width, 2 - channel, y, x); + } + } + } + return values; +} + +} // namespace trtmc diff --git a/families/yolox/runtime/image_preprocess_seam.h b/families/yolox/runtime/image_preprocess_seam.h new file mode 100644 index 0000000000..4cc102fe06 --- /dev/null +++ b/families/yolox/runtime/image_preprocess_seam.h @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace trtmc { + +struct YoloxPreprocessConfig { + std::int32_t input_image_h{640}; + std::int32_t input_image_w{640}; + float pad_value{114.0F}; +}; + +// How the source image was fitted into the square network input. The pipeline +// needs it to map boxes back, so it is returned rather than recomputed. +struct YoloxLetterbox { + float scale{1.0F}; + float pad_x{0.0F}; + float pad_y{0.0F}; +}; + +// `pixels` is an interleaved RGB image in [0, 1], the layout the CLI's +// image reader produces. The result is planar BGR CHW in [0, 255] for the engine. +std::vector preprocess_yolox_image(const float* pixels, std::int32_t height, + std::int32_t width, const YoloxPreprocessConfig& config, + YoloxLetterbox& letterbox); + +} // namespace trtmc diff --git a/families/yolox/runtime/pipeline.cpp b/families/yolox/runtime/pipeline.cpp new file mode 100644 index 0000000000..7b04a599ea --- /dev/null +++ b/families/yolox/runtime/pipeline.cpp @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/yolox/runtime/pipeline.h" + +#include +#include +#include +#include +#include +#include + +namespace trtmc { +namespace { + +const Tensor* find(const TensorMap& outputs, const std::string& name) { + const auto entry = outputs.find(name); + return entry == outputs.end() ? nullptr : &entry->second; +} + +// Intersection over union of two corner-form boxes. +float overlap(const DetectionBox& left, const DetectionBox& right) { + const float x0 = std::max(left.x_min, right.x_min); + const float y0 = std::max(left.y_min, right.y_min); + const float x1 = std::min(left.x_max, right.x_max); + const float y1 = std::min(left.y_max, right.y_max); + const float shared = std::max(0.0F, x1 - x0) * std::max(0.0F, y1 - y0); + if (shared <= 0.0F) + return 0.0F; + const auto area = [](const DetectionBox& box) { + return std::max(0.0F, box.x_max - box.x_min) * std::max(0.0F, box.y_max - box.y_min); + }; + const float total = area(left) + area(right) - shared; + return total > 0.0F ? shared / total : 0.0F; +} + +// Greedy non-maximum suppression, per class. YOLOX's head reports one +// prediction per cell and leaves the overlaps in, so the runtime removes +// them. Boxes of different classes never suppress each other, which is what +// the reference does unless it is asked for the class-agnostic variant. +} // namespace + +std::vector suppress_yolox_boxes(std::vector boxes, float iou_threshold, + std::size_t max_detections) { + if (max_detections == 0) + return {}; + std::stable_sort(boxes.begin(), boxes.end(), [](const DetectionBox& a, const DetectionBox& b) { + return a.score > b.score; + }); + std::vector kept; + std::vector dropped(boxes.size(), false); + for (std::size_t index = 0; index < boxes.size(); ++index) { + if (dropped[index]) + continue; + kept.push_back(boxes[index]); + if (kept.size() >= max_detections) + break; + for (std::size_t other = index + 1; other < boxes.size(); ++other) { + if (dropped[other] || boxes[other].class_id != boxes[index].class_id) + continue; + if (overlap(boxes[index], boxes[other]) > iou_threshold) + dropped[other] = true; + } + } + return kept; +} + +YoloxObjectDetectionPipeline::YoloxObjectDetectionPipeline(std::unique_ptr model, + YoloxPreprocessConfig preprocess_config, + float score_threshold, + float iou_threshold, + std::int32_t max_detections) + : model_(std::move(model)), preprocess_config_(std::move(preprocess_config)), + score_threshold_(score_threshold), iou_threshold_(iou_threshold), + max_detections_(max_detections) { + if (!model_ || !model_->ok()) + throw std::runtime_error("YoloxObjectDetectionPipeline: invalid model"); +} + +ObjectDetectionResult YoloxObjectDetectionPipeline::detect(const float* pixels, int32_t height, + int32_t width) { + YoloxLetterbox letterbox; + auto values = preprocess_yolox_image(pixels, height, width, preprocess_config_, letterbox); + Tensor input; + input.data = values.data(); + input.shape = {1, 3, preprocess_config_.input_image_h, preprocess_config_.input_image_w}; + input.dtype = DType::kFloat32; + const auto outputs = model_->forward({{"pixel_values", input}}); + + const Tensor* boxes = find(outputs, "boxes"); + const Tensor* scores = find(outputs, "scores"); + const Tensor* classes = find(outputs, "classes"); + if (boxes == nullptr || scores == nullptr || classes == nullptr) + throw std::runtime_error("YOLOX engine did not return boxes, scores and classes"); + if (boxes->dtype != DType::kFloat32 || scores->dtype != DType::kFloat32) + throw std::runtime_error("YOLOX boxes and scores must be float32"); + if (classes->dtype != DType::kInt32) + throw std::runtime_error("YOLOX classes must be int32"); + + const auto count = static_cast(scores->numel()); + if (static_cast(boxes->numel()) != count * 4U || + static_cast(classes->numel()) != count) + throw std::runtime_error("YOLOX detection outputs disagree on their length"); + + std::size_t expected_count = 0; + for (const int stride : {8, 16, 32}) + expected_count += static_cast(preprocess_config_.input_image_h / stride) * + (preprocess_config_.input_image_w / stride); + if (count != expected_count || boxes->data == nullptr || scores->data == nullptr || + classes->data == nullptr) + throw std::runtime_error("YOLOX detection slots do not match the configured input size"); + + const auto* box_values = static_cast(boxes->data); + const auto* score_values = static_cast(scores->data); + const auto* class_values = static_cast(classes->data); + + std::vector candidates; + candidates.reserve(count); + for (std::size_t index = 0; index < count; ++index) { + const float score = score_values[index]; + // Every cell is reported, in no particular order, so the whole set + // has to be walked rather than stopped at the first weak one. + if (!(score >= score_threshold_)) + continue; + DetectionBox box; + // Undo the letterbox: remove the padding, then the scale. + const float left = (box_values[index * 4U + 0U] - letterbox.pad_x) / letterbox.scale; + const float top = (box_values[index * 4U + 1U] - letterbox.pad_y) / letterbox.scale; + const float right = (box_values[index * 4U + 2U] - letterbox.pad_x) / letterbox.scale; + const float bottom = (box_values[index * 4U + 3U] - letterbox.pad_y) / letterbox.scale; + // Upstream runs NMS on unbounded corners and reports coordinates / r. + // Clipping before suppression changes the IoU of border detections. + if (!std::isfinite(score) || !std::isfinite(left) || !std::isfinite(top) || + !std::isfinite(right) || !std::isfinite(bottom) || right < left || bottom < top || + class_values[index] < 0 || class_values[index] >= 80) + throw std::runtime_error("YOLOX engine returned an invalid detection"); + box.x_min = left; + box.y_min = top; + box.x_max = right; + box.y_max = bottom; + box.score = score; + box.class_id = class_values[index]; + candidates.push_back(box); + } + + ObjectDetectionResult result; + result.image_height = height; + result.image_width = width; + result.boxes = suppress_yolox_boxes(std::move(candidates), iou_threshold_, + static_cast(max_detections_)); + return result; +} + +} // namespace trtmc diff --git a/families/yolox/runtime/pipeline.h b/families/yolox/runtime/pipeline.h new file mode 100644 index 0000000000..ca56636353 --- /dev/null +++ b/families/yolox/runtime/pipeline.h @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "families/yolox/runtime/image_preprocess_seam.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/task.h" + +#include +#include +#include +#include + +namespace trtmc { + +// Greedy non-maximum suppression, per class, ordered by score. Exposed so the +// suppression can be tested without standing up an engine. +std::vector suppress_yolox_boxes(std::vector boxes, float iou_threshold, + std::size_t max_detections); + +class YoloxObjectDetectionPipeline final : public IObjectDetection { + public: + YoloxObjectDetectionPipeline(std::unique_ptr model, + YoloxPreprocessConfig preprocess_config, float score_threshold, + float iou_threshold, std::int32_t max_detections); + + ObjectDetectionResult detect(const float* pixels, int32_t height, int32_t width) override; + + private: + std::unique_ptr model_; + YoloxPreprocessConfig preprocess_config_; + float score_threshold_; + float iou_threshold_; + std::int32_t max_detections_; +}; + +} // namespace trtmc diff --git a/families/yolox/runtime/plugin.cpp b/families/yolox/runtime/plugin.cpp new file mode 100644 index 0000000000..28634f4d23 --- /dev/null +++ b/families/yolox/runtime/plugin.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/yolox/runtime/pipeline.h" +#include "trtmc/runtime/family_factory.h" +#include "trtmc/runtime/trt_backend.h" + +#include +#include +#include +#include + +namespace trtmc::yolox { +namespace { + +constexpr float kScoreThreshold = 0.25F; +constexpr float kIouThreshold = 0.45F; + +std::int32_t detection_slots(const YoloxPreprocessConfig& config) { + std::int32_t count = 0; + for (const int stride : {8, 16, 32}) + count += (config.input_image_h / stride) * (config.input_image_w / stride); + return count; +} + +std::vector require_section(const BundleReader& bundle, const char* name) { + const auto* section = bundle.find_section(name); + if (section == nullptr || section->length == 0) + throw std::runtime_error("YOLOX bundle section is missing or empty: " + std::string(name)); + return bundle.read_section(name); +} + +YoloxPreprocessConfig parse_config(const std::vector& data) { + const auto json = nlohmann::json::parse(data.begin(), data.end()); + YoloxPreprocessConfig config; + config.input_image_h = json.at("input_image_h").get(); + config.input_image_w = json.at("input_image_w").get(); + config.pad_value = json.at("pad_value").get(); + const auto score = json.at("score_threshold").get(); + const auto iou = json.at("iou_threshold").get(); + const auto maximum = json.at("max_detections").get(); + if ((config.input_image_h != 416 && config.input_image_h != 640) || + config.input_image_w != config.input_image_h || config.pad_value != 114.0F || + json.at("num_classes").get() != 80 || maximum != detection_slots(config) || + score != kScoreThreshold || iou != kIouThreshold) + throw std::runtime_error("YOLOX runtime.json does not match its contract"); + return config; +} + +std::unique_ptr load_engine(IBackend& backend, const std::vector& plan) { + ModuleCreateOptions options{}; + auto engine = backend.create_module(plan.data(), plan.size(), options); + if (!engine || !engine->ok()) + throw std::runtime_error("YOLOX engine failed to load"); + return engine; +} + +} // namespace +} // namespace trtmc::yolox + +extern "C" trtmc::ITask* trtmc_create_family(const trtmc::FamilyContext& context) { + if (context.kv_cache_size_bytes != 0) + throw std::invalid_argument("yolox does not support --kv-cache-size"); + const auto config_data = trtmc::yolox::require_section(context.reader, "runtime.json"); + const auto plan = trtmc::yolox::require_section(context.reader, "engine.plan"); + auto config = trtmc::yolox::parse_config(config_data); + auto engine = trtmc::yolox::load_engine(context.backend, plan); + const auto maximum = trtmc::yolox::detection_slots(config); + return new trtmc::YoloxObjectDetectionPipeline(std::move(engine), std::move(config), + trtmc::yolox::kScoreThreshold, + trtmc::yolox::kIouThreshold, maximum); +} diff --git a/families/yolox/support.py b/families/yolox/support.py new file mode 100644 index 0000000000..d907a59cfe --- /dev/null +++ b/families/yolox/support.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exact local checkpoint identity and public task owned by YOLOX.""" + +from tensorrt_model_connect.model_support import FamilySupport, ModelMetadata + + +ARCHIVES = ( + "yolox_nano.pth", + "yolox_tiny.pth", + "yolox_s.pth", + "yolox_m.pth", + "yolox_l.pth", + "yolox_x.pth", + "yolox_darknet.pth", +) + + +def describe(metadata: ModelMetadata) -> FamilySupport | None: + if any(name in metadata.files for name in ARCHIVES): + return FamilySupport(tasks=("object_detection",), default_task="object_detection") + return None diff --git a/families/yolox/tests/__init__.py b/families/yolox/tests/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/families/yolox/tests/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/families/yolox/tests/cpp/test_image_preprocess_seam.cpp b/families/yolox/tests/cpp/test_image_preprocess_seam.cpp new file mode 100644 index 0000000000..2e44720373 --- /dev/null +++ b/families/yolox/tests/cpp/test_image_preprocess_seam.cpp @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/yolox/runtime/image_preprocess_seam.h" +#include "families/yolox/runtime/pipeline.h" + +#include +#include +#include +#include +#include +#include + +namespace { +void require(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(message); +} +} // namespace + +int main(int argc, char** argv) { + try { + trtmc::YoloxPreprocessConfig config; + trtmc::YoloxLetterbox letterbox; + if (argc == 5 || argc == 6) { + // Family test seam: raw interleaved RGB float32 in, BGR CHW out. + const int height = std::stoi(argv[2]); + const int width = std::stoi(argv[3]); + if (argc == 6) + config.input_image_h = config.input_image_w = std::stoi(argv[5]); + require(height > 0 && width > 0, "invalid image dimensions"); + std::vector pixels(static_cast(height) * width * 3U); + std::ifstream input(argv[1], std::ios::binary); + input.read(reinterpret_cast(pixels.data()), pixels.size() * sizeof(float)); + require(static_cast(input), "could not read input pixels"); + const auto values = + trtmc::preprocess_yolox_image(pixels.data(), height, width, config, letterbox); + std::ofstream output(argv[4], std::ios::binary); + output.write(reinterpret_cast(values.data()), + values.size() * sizeof(float)); + require(static_cast(output), "could not write preprocessed pixels"); + return 0; + } + require(argc == 1, "expected no arguments or input height width output [size]"); + config.input_image_h = config.input_image_w = 2; + const float red_blue[] = {1, 0, 0, 0, 0, 1}; + const auto values = trtmc::preprocess_yolox_image(red_blue, 1, 2, config, letterbox); + const std::vector expected = {0, 255, 114, 114, 0, 0, 114, 114, 255, 0, 114, 114}; + require(values == expected, "BGR bytes, top-left placement or bottom padding is wrong"); + require(letterbox.scale == 1 && letterbox.pad_x == 0 && letterbox.pad_y == 0, + "YOLOX must not center its letterbox"); + config.input_image_h = config.input_image_w = 3; + const float corners[] = {0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1}; + const auto sampled = trtmc::preprocess_yolox_image(corners, 2, 2, config, letterbox); + const std::vector expected_samples = {0, 0, 0, 0, 64, 128, 0, 128, 255, + 0, 0, 0, 128, 64, 0, 255, 128, 0, + 0, 128, 255, 0, 64, 128, 0, 0, 0}; + require(sampled == expected_samples, + "bilinear channel selection, coordinates or byte rounding is wrong"); + config.input_image_h = config.input_image_w = 4; + const std::vector constant(2U * 3U * 3U, 1.0F); + const auto resized = + trtmc::preprocess_yolox_image(constant.data(), 2, 3, config, letterbox); + require(std::fabs(letterbox.scale - 4.0F / 3.0F) < 1e-6F, "resize ratio is wrong"); + for (std::size_t channel = 0; channel < 3; ++channel) { + for (std::size_t index = 0; index < 16; ++index) { + require(resized[channel * 16 + index] == (index < 8 ? 255.0F : 114.0F), + "resized dimensions must truncate, with padding at the bottom"); + } + } + bool rejected = false; + try { + trtmc::preprocess_yolox_image(nullptr, 1, 2, config, letterbox); + } catch (const std::invalid_argument&) { + rejected = true; + } + require(rejected, "empty input was accepted"); + const std::vector boxes = {{-100, 0, 10, 10, 0.9F, 0}, + {0, 0, 10, 10, 0.8F, 0}, + {0, 0, 10, 10, 0.7F, 1}, + {0, 0, 10, 10, 0.6F, 0}}; + const auto kept = trtmc::suppress_yolox_boxes(boxes, 0.45F, 8400); + require(kept.size() == 3, "NMS must preserve unbounded boxes and distinct classes"); + require(kept[0].x_min == -100 && kept[1].score == 0.8F && kept[2].class_id == 1, + "NMS output ordering or coordinates are wrong"); + require(trtmc::suppress_yolox_boxes(boxes, 0.45F, 0).empty(), "zero cap is not empty"); + std::cout << "YOLOX preprocessing and NMS checks passed\n"; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/families/yolox/tests/data/test_img.jpeg b/families/yolox/tests/data/test_img.jpeg new file mode 100644 index 0000000000..095c03620f Binary files /dev/null and b/families/yolox/tests/data/test_img.jpeg differ diff --git a/families/yolox/tests/manifests/yolox-s.json b/families/yolox/tests/manifests/yolox-s.json new file mode 100644 index 0000000000..f530d0584c --- /dev/null +++ b/families/yolox/tests/manifests/yolox-s.json @@ -0,0 +1,22 @@ +{ + "name": "yolox-s", + "family": "yolox", + "task": "object_detection", + "bundle": "yolox-s.bundle", + "precision": "fp16", + "tensor_parallel_size": 1, + "external_files": [ + { + "path": "yolox_s.pth", + "url": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth" + } + ], + "reference_revision": "6ddff4824372906469a7fae2dc3206c7aa4bbaee", + "testcases": [ + { + "name": "yolox-s", + "premerge": true, + "test_image": "data/test_img.jpeg" + } + ] +} diff --git a/families/yolox/tests/reference-source.json b/families/yolox/tests/reference-source.json new file mode 100644 index 0000000000..ff44ddc7b6 --- /dev/null +++ b/families/yolox/tests/reference-source.json @@ -0,0 +1,4 @@ +{ + "repository": "Megvii-BaseDetection/YOLOX", + "revision": "6ddff4824372906469a7fae2dc3206c7aa4bbaee" +} diff --git a/families/yolox/tests/test_e2e.py b/families/yolox/tests/test_e2e.py new file mode 100644 index 0000000000..d8ea3759f4 --- /dev/null +++ b/families/yolox/tests/test_e2e.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct build, native runtime, and official Megvii reference proof for YOLOX.""" + +from __future__ import annotations + +import json +import struct +import sys +import os +import subprocess +from pathlib import Path + +import numpy as np +import pytest + +from tensorrt_model_connect import BuildRequest, build + + +FAMILY = "yolox" +TEST_ROOT = Path(__file__).resolve().parent +MANIFEST_ROOT = TEST_ROOT / "manifests" + + +def _cases() -> dict[str, tuple[dict, dict]]: + result = {} + for path in sorted(MANIFEST_ROOT.glob("*.json")): + manifest = json.loads(path.read_text(encoding="utf-8")) + assert manifest["family"] == FAMILY + assert manifest["task"] == "object_detection" + for case in manifest["testcases"]: + name = str(case["name"]) + assert name not in result + result[name] = (manifest, case) + assert result + return result + + +CASES = _cases() + + +def _selection(config) -> set[str]: + selected = set() + for option in ("--e2e-model", "--e2e-testcase"): + for raw in config.getoption(option, default=[]) or []: + selected.update(value.strip() for value in str(raw).split(",") if value.strip()) + models_file = config.getoption("--e2e-models-file", default=None) + if models_file: + selected.update( + line.strip() + for line in Path(models_file).read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + ) + return selected + + +def pytest_generate_tests(metafunc) -> None: + if "case_name" not in metafunc.fixturenames: + return + selected = _selection(metafunc.config) + names = [ + name + for name, (manifest, _) in CASES.items() + if not selected or selected & {FAMILY, name, manifest["name"]} + ] + if not selected: + names = [ + pytest.param( + name, + marks=pytest.mark.skip(reason="real YOLOX E2E requires explicit selection"), + id=name, + ) + for name in names + ] + metafunc.parametrize("case_name", names) + + +def _required_path(value: str | None, label: str) -> Path: + assert value, f"selected YOLOX E2E requires {label}" + path = Path(value) + assert path.exists(), f"selected YOLOX E2E {label} does not exist: {path}" + return path + + +def _reference_root(manifest): + root = _required_path( + os.environ.get("TRTMC_REFERENCE_SOURCE_DIR"), "TRTMC_REFERENCE_SOURCE_DIR" + ) + revision = subprocess.check_output( + ["git", "-C", str(root), "rev-parse", "HEAD"], text=True + ).strip() + assert revision == manifest["reference_revision"], ( + "reference checkout must match the pinned revision" + ) + dirty = subprocess.check_output( + ["git", "-C", str(root), "diff", "--name-only", "HEAD"], text=True + ) + assert not dirty.strip(), "the official reference must be unmodified" + metadata = json.loads((TEST_ROOT / "reference-source.json").read_text()) + assert metadata["revision"] == manifest["reference_revision"] + sys.path.insert(0, str(root)) + return root + + +def _native_pixels(image, binary, tmp_path, image_size=640): + height, width = image.shape[:2] + rgb = np.ascontiguousarray(image[..., ::-1], dtype=np.float32) / 255.0 + input_path, output_path = tmp_path / "rgb.f32", tmp_path / "preprocessed.f32" + rgb.tofile(input_path) + subprocess.run( + [str(binary), str(input_path), str(height), str(width), str(output_path), str(image_size)], + check=True, + capture_output=True, + timeout=30, + ) + return np.fromfile(output_path, dtype=np.float32).reshape(3, image_size, image_size) + + +def _engine_outputs(bundle, pixels): + import tensorrt as trt + import torch + + # Read the public bundle container to replay the exact engine built above. + with bundle.open("rb") as stream: + assert stream.read(8) == b"BUNDLE\x01\x00" + size = struct.unpack(" 0, "fixture must exercise real detections" + assert len(boxes) == len(expected), (actual, expected.tolist()) + np.testing.assert_array_equal(classes, expected[:, 6].astype(np.int32)) + np.testing.assert_allclose(scores, expected[:, 4] * expected[:, 5], rtol=0, atol=0.01) + target = expected[:, :4] / ratio + # Two pixels in network coordinates, independent of original image size. + np.testing.assert_allclose(boxes * ratio, target * ratio, rtol=0, atol=2.0) + intersection = np.maximum( + 0, np.minimum(boxes[:, 2:], target[:, 2:]) - np.maximum(boxes[:, :2], target[:, :2]) + ).prod(axis=1) + union = ( + (boxes[:, 2:] - boxes[:, :2]).prod(axis=1) + + (target[:, 2:] - target[:, :2]).prod(axis=1) + - intersection + ) + assert np.all(intersection / union >= 0.98), intersection / union + + +def test_official_checkpoint_e2e(case_name: str, tmp_path: Path) -> None: + manifest, case = CASES[case_name] + reference_root = _reference_root(manifest) + import cv2 + import torch + import yolox + from yolox.exp import get_exp + from yolox.data.data_augment import preproc + from yolox.utils import postprocess + + assert Path(yolox.__file__).resolve().is_relative_to(reference_root.resolve()) + binary = _required_path(os.environ.get("TRTMC_BINARY"), "TRTMC_BINARY") + runtime_root = _required_path(os.environ.get("TRTMC_RUNTIME_ROOT"), "TRTMC_RUNTIME_ROOT") + native_build = Path(os.environ.get("TRTMC_NATIVE_BUILD_DIR", str(runtime_root))) + seam = native_build / "families/yolox/test_yolox_image_preprocess" + assert seam.is_file(), f"selected YOLOX E2E requires the native seam test: {seam}" + checkpoints = _required_path(os.environ.get("TRTMC_YOLOX_MODEL_DIR"), "TRTMC_YOLOX_MODEL_DIR") + assert (runtime_root / "libtrtmc_backend_trt.so").is_file() + assert (runtime_root / "libtrtmc_model_yolox.so").is_file() + checkpoint_path = checkpoints / manifest["external_files"][0]["path"] + assert checkpoint_path.is_file(), f"selected YOLOX checkpoint is missing: {checkpoint_path}" + # Each build sees exactly its selected archive, even when CI stages all sizes together. + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / checkpoint_path.name).symlink_to(checkpoint_path.resolve()) + assert torch.cuda.is_available(), "selected YOLOX E2E requires a CUDA GPU" + bundle = tmp_path / manifest["bundle"] + build( + BuildRequest( + model_dir=model_dir, + output_path=bundle, + family=FAMILY, + task=manifest["task"], + precision=manifest["precision"], + tensor_parallel_size=manifest["tensor_parallel_size"], + ) + ) + experiment_name = "yolov3" if checkpoint_path.stem == "yolox_darknet" else checkpoint_path.stem + experiment = get_exp(str(reference_root / f"exps/default/{experiment_name}.py"), None) + reference = experiment.get_model().eval().cuda() + image_size = experiment.test_size[0] + reference.load_state_dict( + torch.load(checkpoint_path, map_location="cpu", weights_only=True)["model"], strict=True + ) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + original = cv2.imread(str(TEST_ROOT / case["test_image"])) + assert original is not None + height, width = original.shape[:2] + portrait = np.full((width * 2, width, 3), 114, dtype=np.uint8) + portrait[:height] = original + for label, image in (("landscape", original), ("portrait", portrait)): + official_pixels, ratio = preproc(image, experiment.test_size) + native_pixels = _native_pixels(image, seam, tmp_path, image_size) + # OpenCV's byte resize uses fixed-point rounding. No channel, padding, + # scale or normalization error can fit within this one-byte bound. + np.testing.assert_allclose(native_pixels, official_pixels, rtol=0, atol=1.0) + actual_raw = _engine_outputs(bundle, native_pixels) + with torch.no_grad(): + same_input = reference(torch.from_numpy(native_pixels[None]).cuda())[0].cpu().numpy() + official = reference(torch.from_numpy(official_pixels[None]).cuda()) + expected = postprocess( + official.clone(), 80, conf_thre=0.25, nms_thre=0.45, class_agnostic=False + )[0] + expected_scores = same_input[:, 4] * same_input[:, 5:].max(axis=1) + expected_classes = same_input[:, 5:].argmax(axis=1) + expected_boxes = np.concatenate( + ( + same_input[:, :2] - same_input[:, 2:4] * 0.5, + same_input[:, :2] + same_input[:, 2:4] * 0.5, + ), + axis=1, + ) + prediction_count = sum((image_size // stride) ** 2 for stride in (8, 16, 32)) + assert actual_raw["boxes"].shape == (prediction_count, 4) + assert all(np.isfinite(value).all() for value in actual_raw.values()) + np.testing.assert_allclose(actual_raw["scores"], expected_scores, rtol=0, atol=0.01) + foreground = (expected_scores >= 0.1) | (actual_raw["scores"] >= 0.1) + assert foreground.any() + np.testing.assert_array_equal( + actual_raw["classes"][foreground], expected_classes[foreground] + ) + np.testing.assert_allclose( + actual_raw["boxes"][foreground], expected_boxes[foreground], rtol=0, atol=2.0 + ) + image_path = tmp_path / f"{label}.png" + assert cv2.imwrite(str(image_path), image) + completed = subprocess.run( + [ + str(binary), + "detect", + str(bundle), + "--runtime-root", + str(runtime_root), + "--image", + str(image_path), + ], + check=True, + capture_output=True, + text=True, + timeout=120, + ) + actual = json.loads(completed.stdout) + assert expected is not None + _compare_detections(actual, expected, ratio) + print(f"{case_name} {label}: {len(expected)} detections match the official reference") diff --git a/families/yolox/tests/test_model.py b/families/yolox/tests/test_model.py new file mode 100644 index 0000000000..bf78094017 --- /dev/null +++ b/families/yolox/tests/test_model.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint, request, BN-folding and TensorRT activation contracts.""" + +from dataclasses import replace +from pathlib import Path + +import numpy as np +import pytest +import tensorrt as trt +import torch + +from families.yolox import graph +from families.yolox.checkpoint import Checkpoint +from families.yolox.model import _fold, build +from families.yolox.support import ARCHIVES, describe +from tensorrt_model_connect import BuildRequest +from tensorrt_model_connect.model_support import ModelMetadata + + +@pytest.mark.parametrize("name", ARCHIVES) +def test_exact_checkpoint_identity(name): + assert describe(ModelMetadata(config={}, model_index={}, files=(name,))) is not None + + +def test_other_checkpoint_names_are_not_claimed(): + for name in ("yolov5n.pt", "yolox.pth", "yolox_custom.pth"): + assert describe(ModelMetadata(config={}, model_index={}, files=(name,))) is None + + +@pytest.mark.parametrize( + "name,size", + [ + ("nano", 416), + ("tiny", 416), + ("s", 640), + ("m", 640), + ("l", 640), + ("x", 640), + ("darknet", 640), + ], +) +def test_checkpoint_selection_and_published_image_size(tmp_path, name, size): + torch.save({"model": {"head.weight": torch.ones(1)}}, tmp_path / f"yolox_{name}.pth") + checkpoint = Checkpoint.open(tmp_path) + assert checkpoint.image_size == size + np.testing.assert_array_equal(checkpoint.tensor("head.weight"), [1.0]) + + +def test_checkpoint_selection_rejects_missing_or_ambiguous_archives(tmp_path): + with pytest.raises(ValueError, match="exactly one"): + Checkpoint.open(tmp_path) + for name in ("yolox_s.pth", "yolox_m.pth"): + torch.save({"model": {"head.weight": torch.ones(1)}}, tmp_path / name) + with pytest.raises(ValueError, match="exactly one"): + Checkpoint.open(tmp_path) + + +@pytest.mark.parametrize( + "field,value", + [ + ("backend", "trt_rtx"), + ("task", "image_classification"), + ("dynamic_kv_cache", True), + ("image_height", 320), + ("image_width", 320), + ("video_num_frames", 2), + ("max_batch_size", 2), + ("tensor_parallel_size", 2), + ("context_parallel_size", 2), + ("quantization", "fp8"), + ("fp32_layers", (0,)), + ("max_sequence_length", 2), + ], +) +def test_unsupported_request_fails_before_checkpoint_access(field, value): + request = BuildRequest( + model_dir=Path("missing"), + output_path=Path("unused.bundle"), + family="yolox", + task="object_detection", + precision="fp16", + ) + with pytest.raises((NotImplementedError, ValueError)): + build(replace(request, **{field: value}), None) + + +def test_fold_matches_pytorch_including_small_variance(): + generator = torch.Generator().manual_seed(10) + conv = torch.nn.Conv2d(3, 4, 3, padding=1, bias=False).eval() + norm = torch.nn.BatchNorm2d(4, eps=1e-3).eval() + with torch.no_grad(): + conv.weight.copy_(torch.randn(conv.weight.shape, generator=generator) * 0.1) + norm.weight.copy_(torch.tensor([0.5, 2.0, -0.1, 1.0])) + norm.bias.copy_(torch.tensor([1.0, -2.0, 0.3, 0.0])) + norm.running_mean.copy_(torch.tensor([0.2, -0.4, 0.0, 1.0])) + norm.running_var.copy_(torch.tensor([0.0001, 0.01, 1.0, 4.0])) + state = { + "block.conv.weight": conv.weight.detach(), + **{f"block.bn.{name}": value for name, value in norm.state_dict().items()}, + } + weight, bias = _fold(Checkpoint(state), "block", np.float32) + pixels = torch.randn((1, 3, 7, 9), generator=generator) + actual = torch.nn.functional.conv2d( + pixels, torch.from_numpy(weight), torch.from_numpy(bias), padding=1 + ) + torch.testing.assert_close(actual, norm(conv(pixels)), atol=2e-5, rtol=2e-5) + # Wrong PyTorch-default epsilon must fail this numerical oracle. + norm.eps = 1e-5 + assert (actual - norm(conv(pixels))).abs().max() > 1.0 + + +def test_checkpoint_rejects_missing_extra_and_nonfinite_tensors(tmp_path): + torch.save({"model": {"head.weight": torch.ones(1)}}, tmp_path / "yolox_s.pth") + checkpoint = Checkpoint.open(tmp_path) + with pytest.raises(ValueError, match="missing"): + checkpoint.tensor("backbone.weight") + with pytest.raises(ValueError, match="Unsupported"): + checkpoint.assert_consumed() + with pytest.raises(ValueError, match="non-finite"): + Checkpoint({"weight": torch.tensor([float("nan")])}) + + +@pytest.mark.trt +@pytest.mark.skipif(not torch.cuda.is_available(), reason="TensorRT activation test requires CUDA") +def test_half_silu_matches_the_official_pytorch_activation(): + values = torch.tensor( + [8.0078125, -8.0078125, 1, -1, 0, 3, 7], device="cuda", dtype=torch.float16 + ) + expected = torch.nn.functional.silu(values) + # The old two-operation FP16 expression rounds the positive probe to 8.0. + assert not torch.equal(values * values.sigmoid(), expected) + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + tensor = network.add_input("input", trt.float16, tuple(values.shape)) + output = graph.silu(network, tensor) + output.name = "output" + network.mark_output(output) + plan = builder.build_serialized_network(network, config) + assert plan is not None + runtime = trt.Runtime(logger) + engine = runtime.deserialize_cuda_engine(plan) + assert engine is not None + context = engine.create_execution_context() + actual = torch.empty_like(values) + assert context.set_tensor_address("input", values.data_ptr()) + assert context.set_tensor_address("output", actual.data_ptr()) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + assert context.execute_async_v3(stream.cuda_stream) + stream.synchronize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0)