From 5188b645b1fc680ef8e00d9e2ac442bb6415736e Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 12 Aug 2026 18:45:32 +0000 Subject: [PATCH 1/3] Add PETR ONNX PTQ example Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 1 + examples/onnx_ptq/README.md | 4 + examples/onnx_ptq/petr/Dockerfile | 37 +++ examples/onnx_ptq/petr/README.md | 240 +++++++++++++++++ examples/onnx_ptq/petr/convert_to_fp16.py | 43 +++ examples/onnx_ptq/petr/evaluate.py | 250 ++++++++++++++++++ examples/onnx_ptq/petr/prepare_calibration.py | 105 ++++++++ .../onnx_ptq/petr/prepare_sweep_metadata.py | 164 ++++++++++++ examples/onnx_ptq/petr/quantize.py | 131 +++++++++ .../onnx_ptq/petr/requirements-mmdet3d.txt | 1 + examples/onnx_ptq/petr/requirements-torch.txt | 4 + examples/onnx_ptq/petr/requirements.txt | 25 ++ 12 files changed, 1005 insertions(+) create mode 100644 examples/onnx_ptq/petr/Dockerfile create mode 100644 examples/onnx_ptq/petr/README.md create mode 100644 examples/onnx_ptq/petr/convert_to_fp16.py create mode 100644 examples/onnx_ptq/petr/evaluate.py create mode 100644 examples/onnx_ptq/petr/prepare_calibration.py create mode 100644 examples/onnx_ptq/petr/prepare_sweep_metadata.py create mode 100644 examples/onnx_ptq/petr/quantize.py create mode 100644 examples/onnx_ptq/petr/requirements-mmdet3d.txt create mode 100644 examples/onnx_ptq/petr/requirements-torch.txt create mode 100644 examples/onnx_ptq/petr/requirements.txt diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 776f18178eb..e9c35223997 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,7 @@ Changelog **New Features** +- Add an end-to-end PETRv1 and PETRv2 ONNX PTQ example with nuScenes calibration and accuracy evaluation, INT8 and FP8 backbone/head quantization, and TensorRT engine benchmarking. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers. Weight scales are an orthogonal axis, selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian`` with that algorithm's own options), so one recipe can combine a weight calibration with this activation policy in a single pass. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. diff --git a/examples/onnx_ptq/README.md b/examples/onnx_ptq/README.md index 3cee5535a84..bfabb023d45 100644 --- a/examples/onnx_ptq/README.md +++ b/examples/onnx_ptq/README.md @@ -133,6 +133,10 @@ Inference latency of the model is ms The [FAR3D example](./far3d/) demonstrates an end-to-end workflow that exports and quantizes the FAR3D ONNX image encoder, builds TensorRT engines, and evaluates 3D object detection mAP on the Argoverse 2 validation set. +### PETR 3D object detection + +The [PETR example](./petr/) demonstrates an end-to-end workflow that exports and quantizes PETRv1 and PETRv2 ONNX backbones and heads, builds TensorRT engines, and evaluates 3D object detection mAP on the nuScenes validation set. + ## Advanced Features ### Per node calibration of ONNX models diff --git a/examples/onnx_ptq/petr/Dockerfile b/examples/onnx_ptq/petr/Dockerfile new file mode 100644 index 00000000000..37333a86f7c --- /dev/null +++ b/examples/onnx_ptq/petr/Dockerfile @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +FROM nvcr.io/nvidia/pytorch:26.07-py3 + +ENV LD_LIBRARY_PATH=/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + libgl1 \ + libglib2.0-0 && \ + rm -rf /var/lib/apt/lists/* + +ENV UV_PYTHON_INSTALL_DIR=/opt/python + +RUN python -m pip install --no-cache-dir uv && \ + uv python install 3.8 && \ + uv venv --seed --python 3.8 /opt/petr + +COPY examples/onnx_ptq/petr/requirements*.txt /tmp/petr-requirements/ +RUN env -u PIP_CONSTRAINT /opt/petr/bin/python -m pip install --no-cache-dir \ + -r /tmp/petr-requirements/requirements-torch.txt && \ + env -u PIP_CONSTRAINT /opt/petr/bin/python -m pip install --no-cache-dir \ + -r /tmp/petr-requirements/requirements.txt && \ + env -u PIP_CONSTRAINT /opt/petr/bin/python -m pip install --no-cache-dir \ + --no-build-isolation \ + -r /tmp/petr-requirements/requirements-mmdet3d.txt && \ + mkdir -p /opt/petr/lib/python3.8/site-packages/tensorrt && \ + cp /opt/petr/lib/python3.8/site-packages/tensorrt_bindings/__init__.py \ + /opt/petr/lib/python3.8/site-packages/tensorrt/__init__.py && \ + cp /opt/petr/lib/python3.8/site-packages/tensorrt_bindings/tensorrt.so \ + /opt/petr/lib/python3.8/site-packages/tensorrt/tensorrt.so + +COPY . /opt/Model-Optimizer +RUN cd /opt/Model-Optimizer && \ + env -u PIP_CONSTRAINT python -m pip install --no-cache-dir \ + -e ".[onnx]" \ + "onnxruntime-gpu[cuda,cudnn]~=1.24.2" diff --git a/examples/onnx_ptq/petr/README.md b/examples/onnx_ptq/petr/README.md new file mode 100644 index 00000000000..3575f4b844c --- /dev/null +++ b/examples/onnx_ptq/petr/README.md @@ -0,0 +1,240 @@ +# PETR ONNX PTQ and nuScenes evaluation + +This example quantizes the PETRv1 and PETRv2 image backbones and detection heads to INT8 or FP8 with Model Optimizer. It follows the [NVIDIA DL4AGX PETR TensorRT workflow](https://github.com/NVIDIA/DL4AGX/tree/master/AV-Solutions/petr-trt) and evaluates TensorRT engines on the nuScenes validation set. + +PETR uses a legacy PyTorch/MMCV environment that is incompatible with current Model Optimizer dependencies. The provided image uses `nvcr.io/nvidia/pytorch:26.07-py3` with TensorRT 11.1 and isolates PETR in a Python 3.8 virtual environment. + +## 1. Prepare PETR and nuScenes + +Clone the reference repositories and apply the DL4AGX compatibility patch: + +```bash +git clone https://github.com/NVIDIA/DL4AGX.git +git clone https://github.com/megvii-research/PETR.git +cd PETR +git apply ../DL4AGX/AV-Solutions/petr-trt/patch.diff +git clone https://github.com/open-mmlab/mmdetection3d.git -b v0.17.1 +mkdir -p ckpts data +ln -s /path/to/nuscenes data/nuscenes +cd .. +``` + +Download the `PETR-vov-p4-800x320_epoch24.pth` and `PETRv2-vov-p4-800x320_epoch24.pth` checkpoints linked from the DL4AGX README, rename them as shown below, and prepare this layout: + +```text +PETR/ +├── ckpts/ +│ ├── PETR-vov-p4-800x320_e24.pth +│ └── PETRv2-vov-p4-800x320_e24.pth +├── data/nuscenes/ +│ ├── samples/ +│ ├── sweeps/ +│ ├── v1.0-trainval/ +│ └── nuscenes_infos_val.pkl +└── mmdetection3d/ +``` + +PETRv2 also needs metadata for the previous camera sweeps: + +```bash +/opt/petr/bin/python \ + /opt/Model-Optimizer/examples/onnx_ptq/petr/prepare_sweep_metadata.py \ + /workspace/PETR/data/nuscenes \ + --split val +``` + +This creates `mmdet3d_nuscenes_30f_infos_val.pkl`. + +Build and start the example image from the Model Optimizer checkout: + +```bash +docker build \ + -f examples/onnx_ptq/petr/Dockerfile \ + -t petr-modelopt \ + . + +docker run --rm -it --network=host --gpus=all --shm-size=64G \ + -v /path/to/Model-Optimizer:/opt/Model-Optimizer \ + -v /path/to/PETR:/workspace/PETR \ + -v /path/to/DL4AGX:/workspace/DL4AGX \ + petr-modelopt +``` + +Use the isolated Python 3.8 environment for PETR export, calibration preparation, and evaluation. Keep its site-packages first so the installed MMDetection3D 1.0.0rc6 package takes precedence over the v0.17.1 source tree used by the configs: + +```bash +export PYTHONPATH=/opt/petr/lib/python3.8/site-packages:/workspace/PETR:/workspace/DL4AGX/AV-Solutions/petr-trt/export_eval +``` + +## 2. Export PETRv1 and PETRv2 to ONNX + +Follow the DL4AGX exporter setup: + +```bash +cd /workspace/DL4AGX/AV-Solutions/petr-trt/export_eval +ln -s /workspace/PETR/data data +mkdir -p onnx_files engines + +/opt/petr/bin/python v1/v1_export_to_onnx.py \ + /workspace/PETR/projects/configs/petr/petr_vovnet_gridmask_p4_800x320.py \ + /workspace/PETR/ckpts/PETR-vov-p4-800x320_e24.pth \ + --eval bbox + +/opt/petr/bin/python v2/v2_export_to_onnx.py \ + /workspace/PETR/projects/configs/petrv2/petrv2_vovnet_gridmask_p4_800x320.py \ + /workspace/PETR/ckpts/PETRv2-vov-p4-800x320_e24.pth \ + --eval bbox +``` + +The exporters create a backbone graph and a head graph. Simplify both graphs: + +```bash +for version in v1 v2; do + model="PETR${version}" + python -m onnxsim \ + "onnx_files/${model}.extract_feat.onnx" \ + "onnx_files/sim_${model}.extract_feat.onnx" + python -m onnxsim \ + "onnx_files/${model}.pts_bbox_head.forward.onnx" \ + "onnx_files/sim_${model}.pts_bbox_head.forward.onnx" +done +``` + +## 3. Prepare calibration batches + +Build temporary FP32 engines for collecting the backbone and head inputs: + +```bash +for version in v1 v2; do + model="PETR${version}" + trtexec \ + --onnx="onnx_files/sim_${model}.extract_feat.onnx" \ + --saveEngine="engines/${model}.backbone.calibration.engine" \ + --skipInference + trtexec \ + --onnx="onnx_files/sim_${model}.pts_bbox_head.forward.onnx" \ + --saveEngine="engines/${model}.head.calibration.engine" \ + --skipInference +done +``` + +Collect 512 representative backbone and head input batches. The default interval samples across the 6,019-frame validation set: + +```bash +/opt/petr/bin/python /opt/Model-Optimizer/examples/onnx_ptq/petr/prepare_calibration.py \ + v1 \ + /workspace/PETR/projects/configs/petr/petr_vovnet_gridmask_p4_800x320.py \ + /workspace/PETR/ckpts/PETR-vov-p4-800x320_e24.pth \ + onnx_files/sim_PETRv1.extract_feat.onnx \ + onnx_files/sim_PETRv1.pts_bbox_head.forward.onnx \ + engines/PETRv1.backbone.calibration.engine \ + engines/PETRv1.head.calibration.engine \ + calibration/PETRv1 +``` + +Repeat with `v2`, the PETRv2 config, checkpoint, graphs, engines, and `calibration/PETRv2` output directory. + +TensorRT 11 uses strongly typed networks and no longer accepts `--fp16`. +Create calibrated mixed-FP16 graphs with Model Optimizer AutoCast, preserving +FP32 model inputs and outputs, then build the reference engines. The example +below uses one representative batch for AutoCast node classification: + +```bash +for version in v1 v2; do + model="PETR${version}" + python /opt/Model-Optimizer/examples/onnx_ptq/petr/convert_to_fp16.py \ + "onnx_files/sim_${model}.extract_feat.onnx" \ + "onnx_files/sim_${model}.extract_feat.fp16.onnx" \ + "calibration/${model}/backbone/batch_0000.npz" + python /opt/Model-Optimizer/examples/onnx_ptq/petr/convert_to_fp16.py \ + "onnx_files/sim_${model}.pts_bbox_head.forward.onnx" \ + "onnx_files/sim_${model}.pts_bbox_head.forward.fp16.onnx" \ + "calibration/${model}/head/batch_0000.npz" + trtexec \ + --onnx="onnx_files/sim_${model}.extract_feat.fp16.onnx" \ + --saveEngine="engines/${model}.backbone.fp16.engine" \ + --skipInference + trtexec \ + --onnx="onnx_files/sim_${model}.pts_bbox_head.forward.fp16.onnx" \ + --saveEngine="engines/${model}.head.fp16.engine" \ + --skipInference +done +``` + +## 4. Quantize the ONNX models + +Use the base Python environment for Model Optimizer. This command quantizes the backbone and keeps the head in FP16: + +```bash +python /opt/Model-Optimizer/examples/onnx_ptq/petr/quantize.py \ + --backbone-onnx onnx_files/sim_PETRv1.extract_feat.onnx \ + --head-onnx onnx_files/sim_PETRv1.pts_bbox_head.forward.onnx \ + --calibration-dir calibration/PETRv1 \ + --precision int8 +``` + +Use `--precision fp8` for FP8. Add `--quantize-head` to quantize both the backbone and head. Repeat the commands for PETRv2. FP8 deployment requires an FP8-capable GPU. + +The backbone quantizer preserves the accuracy-sensitive final VoVNet stage and FPN output layers in FP16, matching the exclusions used by the FAR3D example. + +## 5. Build and evaluate TensorRT engines + +Build quantized graphs as strongly typed engines: + +```bash +precision=int8 +model=PETRv1 +trtexec \ + --onnx="onnx_files/sim_${model}.extract_feat.${precision}.onnx" \ + --saveEngine="engines/${model}.backbone.${precision}.engine" \ + --stronglyTyped \ + --skipInference +trtexec \ + --onnx="onnx_files/sim_${model}.pts_bbox_head.forward.${precision}.onnx" \ + --saveEngine="engines/${model}.head.${precision}.engine" \ + --stronglyTyped \ + --skipInference +``` + +Evaluate any backbone/head pairing. For example, INT8 backbone with FP16 head: + +```bash +/opt/petr/bin/python /opt/Model-Optimizer/examples/onnx_ptq/petr/evaluate.py \ + v1 \ + /workspace/PETR/projects/configs/petr/petr_vovnet_gridmask_p4_800x320.py \ + /workspace/PETR/ckpts/PETR-vov-p4-800x320_e24.pth \ + engines/PETRv1.backbone.int8.engine \ + engines/PETRv1.head.fp16.engine +``` + +Use `--max-samples N` for an inference smoke test. Dataset metrics are skipped when only part of the validation set is processed. + +Measure each engine with host/device transfers disabled and add the backbone and head median GPU compute times: + +```bash +trtexec --loadEngine=engines/PETRv1.backbone.int8.engine \ + --noDataTransfers --useCudaGraph --warmUp=1000 --duration=10 +trtexec --loadEngine=engines/PETRv1.head.fp16.engine \ + --noDataTransfers --useCudaGraph --warmUp=1000 --duration=10 +``` + +## Results on the nuScenes validation set + +Results below use TensorRT 11.1.0.106 on an NVIDIA RTX 6000 Ada Generation GPU. Accuracy is measured over all 6,019 validation samples after calibration with 512 batches. GPU compute time is the sum of the backbone and head median times reported by `trtexec`; it excludes data transfers and PETRv2's reusable previous-frame feature extraction. + +| Model | Backbone precision | Head precision | Framework | GPU compute time (ms) | Accuracy (mAP) | +| --- | --- | --- | --- | ---: | ---: | +| PETRv1-vov-p4-800x320 | FP16 | FP16 | TensorRT 11.1 | 14.507 | 0.3781 | +| PETRv1-vov-p4-800x320 | INT8 | FP16 | TensorRT 11.1 | 9.992 | 0.3711 | +| PETRv1-vov-p4-800x320 | FP8 | FP16 | TensorRT 11.1 | 11.455 | 0.3757 | +| PETRv1-vov-p4-800x320 | INT8 | INT8 | TensorRT 11.1 | 5.877 | 0.1779 | +| PETRv1-vov-p4-800x320 | FP8 | FP8 | TensorRT 11.1 | 8.749 | 0.2071 | +| PETRv2-vov-p4-800x320 | FP16 | FP16 | TensorRT 11.1 | 19.349 | 0.4105 | +| PETRv2-vov-p4-800x320 | INT8 | FP16 | TensorRT 11.1 | 14.468 | 0.4017 | +| PETRv2-vov-p4-800x320 | FP8 | FP16 | TensorRT 11.1 | 16.242 | 0.4092 | +| PETRv2-vov-p4-800x320 | INT8 | INT8 | TensorRT 11.1 | 6.724 | 0.2093 | +| PETRv2-vov-p4-800x320 | FP8 | FP8 | TensorRT 11.1 | 12.261 | 0.2568 | + +TensorRT engines are specific to the TensorRT version and GPU architecture used to build them. These x86 results are not directly comparable with the DRIVE Orin measurements in the DL4AGX reference. + +Quantizing only the backbone retains substantially more accuracy than quantizing the detection head. Fully quantized head results are included to show the measured tradeoff rather than as recommended configurations. diff --git a/examples/onnx_ptq/petr/convert_to_fp16.py b/examples/onnx_ptq/petr/convert_to_fp16.py new file mode 100644 index 00000000000..d2fc0c5aa88 --- /dev/null +++ b/examples/onnx_ptq/petr/convert_to_fp16.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse + +from modelopt.onnx import utils as onnx_utils +from modelopt.onnx.autocast import convert_to_mixed_precision + + +def parse_args(): + parser = argparse.ArgumentParser(description="Convert a PETR ONNX model to FP16") + parser.add_argument("input") + parser.add_argument("output") + parser.add_argument("calibration_data") + return parser.parse_args() + + +def main(): + args = parse_args() + model = convert_to_mixed_precision( + onnx_path=args.input, + low_precision_type="fp16", + calibration_data=args.calibration_data, + keep_io_types=True, + providers=["cuda:0", "cpu"], + ) + onnx_utils.save_onnx(model, args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/onnx_ptq/petr/evaluate.py b/examples/onnx_ptq/petr/evaluate.py new file mode 100644 index 00000000000..31ed8926ecd --- /dev/null +++ b/examples/onnx_ptq/petr/evaluate.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Adapted from https://github.com/NVIDIA/DL4AGX/tree/9f7b29104c253d5bc68334e7b83b3eecb72d4572/AV-Solutions/petr-trt/export_eval. +# Copyright (c) OpenMMLab. All rights reserved. + +import argparse +import importlib +import os + +import tensorrt as trt +import torch +import torch.nn.functional as F +from mmcv import Config, DictAction +from mmcv.runner import load_checkpoint, wrap_fp16_model +from mmcv.utils import import_modules_from_strings +from mmdet.apis import set_random_seed +from mmdet3d.core import bbox3d2result +from mmdet3d.datasets import build_dataloader, build_dataset +from mmdet3d.models import build_model +from tqdm import tqdm + +TRT_TO_TORCH = { + trt.DataType.FLOAT: torch.float32, + trt.DataType.HALF: torch.float16, + trt.DataType.INT8: torch.int8, + trt.DataType.INT32: torch.int32, + trt.DataType.BOOL: torch.bool, + trt.DataType.UINT8: torch.uint8, +} +if int(trt.__version__.split(".")[0]) >= 10: + TRT_TO_TORCH[trt.DataType.INT64] = torch.int64 + +TRT_LOGGER = trt.Logger(trt.Logger.WARNING) +trt.init_libnvinfer_plugins(TRT_LOGGER, "") + + +def aligned_tensor(shape, dtype, device, alignment=256): + element_size = torch.empty((), dtype=dtype).element_size() + element_count = int(torch.tensor(shape).prod().item()) + storage = torch.empty(element_count + alignment // element_size, dtype=dtype, device=device) + offset = ((-storage.data_ptr()) % alignment) // element_size + return storage[offset : offset + element_count].view(shape) + + +class TensorRTRunner: + def __init__(self, engine_path, input_callback=None): + with open(engine_path, "rb") as engine_file: + engine_bytes = engine_file.read() + self.engine = trt.Runtime(TRT_LOGGER).deserialize_cuda_engine(engine_bytes) + if self.engine is None: + raise RuntimeError(f"Failed to deserialize {engine_path}") + self.context = self.engine.create_execution_context() + self.input_callback = input_callback + self.input_names = [] + self.output_names = [] + self.shapes = {} + self.dtypes = {} + for index in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(index) + self.shapes[name] = tuple(self.engine.get_tensor_shape(name)) + self.dtypes[name] = TRT_TO_TORCH[self.engine.get_tensor_dtype(name)] + names = ( + self.input_names + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT + else self.output_names + ) + names.append(name) + + @staticmethod + def fit_shape(value, shape): + if tuple(value.shape) == shape: + return value + if tuple(value.shape[1:]) == shape: + return value.squeeze(0) + if tuple(shape[1:]) == tuple(value.shape): + return value.unsqueeze(0) + raise ValueError(f"Input has shape {tuple(value.shape)}, expected {shape}") + + def __call__(self, stream, values): + if len(values) != len(self.input_names): + raise ValueError(f"Received {len(values)} inputs, expected {len(self.input_names)}") + inputs = [] + for name, value in zip(self.input_names, values): + value = value.to(device="cuda", dtype=self.dtypes[name]).contiguous() + value = self.fit_shape(value, self.shapes[name]) + buffer = aligned_tensor(self.shapes[name], value.dtype, value.device) + buffer.copy_(value) + inputs.append(buffer) + self.context.set_tensor_address(name, buffer.data_ptr()) + if self.input_callback: + self.input_callback(inputs) + + outputs = [] + for name in self.output_names: + output = aligned_tensor(self.shapes[name], self.dtypes[name], "cuda") + outputs.append(output) + self.context.set_tensor_address(name, output.data_ptr()) + if not self.context.execute_async_v3(stream.cuda_stream): + raise RuntimeError("TensorRT execution failed") + stream.synchronize() + return outputs + + +def import_plugin(cfg): + if cfg.get("custom_imports"): + import_modules_from_strings(**cfg.custom_imports) + if cfg.get("plugin"): + plugin_dir = os.path.dirname(cfg.plugin_dir).split("/") + importlib.import_module(".".join(plugin_dir)) + + +def build_runtime(config_path, checkpoint_path, cfg_options=None): + cfg = Config.fromfile(config_path) + if cfg_options: + cfg.merge_from_dict(cfg_options) + import_plugin(cfg) + cfg.model.pretrained = None + cfg.model.train_cfg = None + cfg.data.test.test_mode = True + dataset = build_dataset(cfg.data.test) + loader = build_dataloader( + dataset, + samples_per_gpu=1, + workers_per_gpu=cfg.data.workers_per_gpu, + dist=False, + shuffle=False, + ) + model = build_model(cfg.model, test_cfg=cfg.get("test_cfg")) + if cfg.get("fp16"): + wrap_fp16_model(model) + checkpoint = load_checkpoint(model, checkpoint_path, map_location="cpu") + model.CLASSES = checkpoint.get("meta", {}).get("CLASSES", dataset.CLASSES) + if hasattr(dataset, "PALETTE"): + model.PALETTE = checkpoint.get("meta", {}).get("PALETTE", dataset.PALETTE) + model = model.cuda().eval() + return cfg, dataset, loader, model + + +class PETRPipeline: + def __init__(self, version, model, backbone_engine, head_engine, callbacks=(None, None)): + self.version = version + self.model = model + self.backbone = TensorRTRunner(backbone_engine, callbacks[0]) + self.head = TensorRTRunner(head_engine, callbacks[1]) + + @staticmethod + def masks(features, img_metas): + batch_size, num_cams = features[0].shape[:2] + input_h, input_w, _ = img_metas[0]["pad_shape"][0] + masks = features[0].new_ones((batch_size, num_cams, input_h, input_w)) + for image_id in range(batch_size): + for camera_id in range(num_cams): + image_h, image_w, _ = img_metas[image_id]["img_shape"][camera_id] + masks[image_id, camera_id, :image_h, :image_w] = 0 + return F.interpolate(masks, size=features[0].shape[-2:]).to(torch.bool) + + def backbone_inputs(self, images, img_metas): + if self.version == "v1": + return [images] + current = images[:, :6].contiguous() + previous = images[:, 6:12].contiguous() + previous_features = self.model.extract_img_feat(previous, img_metas) + return [current, *previous_features] + + def head_inputs(self, features, img_metas): + masks = self.masks(features, img_metas) + coords, _ = self.model.pts_bbox_head.position_embeding(features, img_metas, masks) + values = [features[0]] + if self.version == "v2": + timestamps = features[0].new_tensor([meta["timestamp"] for meta in img_metas]) + timestamps = timestamps.view(1, -1, 6) + values.append((timestamps[:, 1] - timestamps[:, 0]).mean(-1)) + values.append(coords) + return values + + def __call__(self, stream, data): + images = data["img"][0].data[0].cuda() + img_metas = data["img_metas"][0].data[0] + with torch.cuda.stream(stream), torch.no_grad(): + features = self.backbone(stream, self.backbone_inputs(images, img_metas)) + camera_count = 6 if self.version == "v1" else 12 + features = [value.reshape(1, camera_count, *value.shape[-3:]) for value in features] + outputs = self.head(stream, self.head_inputs(features, img_metas)) + outputs = [value.float() for value in outputs] + head_outputs = { + "all_cls_scores": outputs[0], + "all_bbox_preds": outputs[1], + "enc_cls_scores": None, + "enc_bbox_preds": None, + } + boxes = self.model.pts_bbox_head.get_bboxes(head_outputs, img_metas, rescale=True) + return [ + {"pts_bbox": bbox3d2result(boxes_3d, scores_3d, labels_3d)} + for boxes_3d, scores_3d, labels_3d in boxes + ] + + +def parse_args(): + parser = argparse.ArgumentParser(description="Evaluate PETR TensorRT engines on nuScenes") + parser.add_argument("version", choices=("v1", "v2")) + parser.add_argument("config") + parser.add_argument("checkpoint") + parser.add_argument("backbone_engine") + parser.add_argument("head_engine") + parser.add_argument("--cfg-options", nargs="+", action=DictAction) + parser.add_argument("--eval-options", nargs="+", action=DictAction) + parser.add_argument("--max-samples", type=int) + args = parser.parse_args() + if args.max_samples is not None and args.max_samples < 1: + raise ValueError("--max-samples must be positive") + return args + + +def main(): + args = parse_args() + set_random_seed(0, deterministic=False) + cfg, dataset, loader, model = build_runtime(args.config, args.checkpoint, args.cfg_options) + pipeline = PETRPipeline(args.version, model, args.backbone_engine, args.head_engine) + stream = torch.cuda.Stream() + outputs = [] + for data in tqdm(loader): + outputs.extend(pipeline(stream, data)) + if args.max_samples is not None and len(outputs) == args.max_samples: + break + if len(outputs) < len(dataset): + print(f"Processed {len(outputs)} samples; skipping dataset metrics") + return + eval_kwargs = cfg.get("evaluation", {}).copy() + for key in ("interval", "tmpdir", "start", "gpu_collect", "save_best", "rule"): + eval_kwargs.pop(key, None) + if args.eval_options: + eval_kwargs.update(args.eval_options) + print(dataset.evaluate(outputs, **eval_kwargs)) + + +if __name__ == "__main__": + main() diff --git a/examples/onnx_ptq/petr/prepare_calibration.py b/examples/onnx_ptq/petr/prepare_calibration.py new file mode 100644 index 00000000000..7d8d00c9911 --- /dev/null +++ b/examples/onnx_ptq/petr/prepare_calibration.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +from pathlib import Path + +import numpy as np +import onnx +import torch +from evaluate import PETRPipeline, build_runtime +from mmdet3d.datasets import build_dataloader +from torch.utils.data import Subset + + +class CalibrationWriter: + def __init__(self, output_dir, onnx_path): + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + if any(self.output_dir.glob("*.npz")): + raise FileExistsError(f"{self.output_dir} already contains calibration batches") + graph = onnx.load(onnx_path, load_external_data=False).graph + self.names = [value.name for value in graph.input] + self.dtypes = [ + onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) + for value in graph.input + ] + self.saved = 0 + + def __call__(self, values): + if len(values) != len(self.names): + raise ValueError(f"Received {len(values)} inputs, expected {len(self.names)}") + batch = { + name: value.detach().cpu().numpy().astype(dtype, copy=False) + for name, dtype, value in zip(self.names, self.dtypes, values) + } + np.savez(self.output_dir / f"batch_{self.saved:04d}.npz", **batch) + self.saved += 1 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Prepare PETR ONNX calibration batches") + parser.add_argument("version", choices=("v1", "v2")) + parser.add_argument("config") + parser.add_argument("checkpoint") + parser.add_argument("backbone_onnx") + parser.add_argument("head_onnx") + parser.add_argument("backbone_engine") + parser.add_argument("head_engine") + parser.add_argument("output_dir", type=Path) + parser.add_argument("--num-samples", type=int, default=512) + parser.add_argument("--sample-skip-interval", type=int, default=10) + return parser.parse_args() + + +def main(): + args = parse_args() + if args.num_samples < 1 or args.sample_skip_interval < 1: + raise ValueError("Sample count and skip interval must be positive") + cfg, dataset, _, model = build_runtime(args.config, args.checkpoint) + stop = min(len(dataset), args.num_samples * args.sample_skip_interval) + subset = Subset(dataset, range(args.sample_skip_interval - 1, stop, args.sample_skip_interval)) + loader = build_dataloader( + subset, + samples_per_gpu=1, + workers_per_gpu=cfg.data.workers_per_gpu, + dist=False, + shuffle=False, + ) + backbone_writer = CalibrationWriter(args.output_dir / "backbone", args.backbone_onnx) + head_writer = CalibrationWriter(args.output_dir / "head", args.head_onnx) + pipeline = PETRPipeline( + args.version, + model, + args.backbone_engine, + args.head_engine, + (backbone_writer, head_writer), + ) + stream = torch.cuda.Stream() + for data in loader: + pipeline(stream, data) + if backbone_writer.saved == args.num_samples: + break + if backbone_writer.saved != args.num_samples or head_writer.saved != args.num_samples: + raise RuntimeError( + f"Prepared {backbone_writer.saved} backbone and {head_writer.saved} head batches; " + f"expected {args.num_samples}" + ) + print(f"Saved {args.num_samples} calibration batches to {args.output_dir}") + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("fork") + main() diff --git a/examples/onnx_ptq/petr/prepare_sweep_metadata.py b/examples/onnx_ptq/petr/prepare_sweep_metadata.py new file mode 100644 index 00000000000..8a01fcd36c4 --- /dev/null +++ b/examples/onnx_ptq/petr/prepare_sweep_metadata.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Adapted from PETR tools/generate_sweep_pkl.py. +# Copyright (c) 2022 Megvii Inc. All rights reserved. + +import argparse +import os +import pickle + +import mmcv +import numpy as np +import tqdm +from nuscenes import NuScenes +from pyquaternion import Quaternion + +sensors = [ + "CAM_FRONT", + "CAM_FRONT_RIGHT", + "CAM_BACK_RIGHT", + "CAM_BACK", + "CAM_BACK_LEFT", + "CAM_FRONT_LEFT", +] + + +def parse_args(): + parser = argparse.ArgumentParser(description="Add PETRv2 camera sweeps to nuScenes metadata") + parser.add_argument("data_root") + parser.add_argument("--split", choices=("train", "val", "test"), default="val") + parser.add_argument("--output") + return parser.parse_args() + + +args = parse_args() +info_prefix = args.split +data_root = os.path.abspath(args.data_root) + os.sep +num_prev = 5 +num_sweep = 5 + +info_path = args.output or os.path.join(data_root, f"mmdet3d_nuscenes_30f_infos_{info_prefix}.pkl") +key_infos = pickle.load(open(os.path.join(data_root, f"nuscenes_infos_{info_prefix}.pkl"), "rb")) +if info_prefix == "test": + nuscenes_version = "v1.0-test" +else: + nuscenes_version = "v1.0-trainval" +nuscenes = NuScenes(nuscenes_version, data_root) + + +def add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat): + sweep_cam = {} + sweep_cam["is_key_frame"] = sample_data["is_key_frame"] + sweep_cam["data_path"] = os.path.join(data_root, sample_data["filename"]) + sweep_cam["type"] = "camera" + sweep_cam["timestamp"] = sample_data["timestamp"] + sweep_cam["sample_data_token"] = sample_data["sample_token"] + pose_record = nuscenes.get("ego_pose", sample_data["ego_pose_token"]) + calibrated_sensor_record = nuscenes.get( + "calibrated_sensor", sample_data["calibrated_sensor_token"] + ) + + sweep_cam["ego2global_translation"] = pose_record["translation"] + sweep_cam["ego2global_rotation"] = pose_record["rotation"] + sweep_cam["sensor2ego_translation"] = calibrated_sensor_record["translation"] + sweep_cam["sensor2ego_rotation"] = calibrated_sensor_record["rotation"] + sweep_cam["cam_intrinsic"] = calibrated_sensor_record["camera_intrinsic"] + + l2e_r_s = sweep_cam["sensor2ego_rotation"] + l2e_t_s = sweep_cam["sensor2ego_translation"] + e2g_r_s = sweep_cam["ego2global_rotation"] + e2g_t_s = sweep_cam["ego2global_translation"] + + l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix + e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix + rotation = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ ( + np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T + ) + translation = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ ( + np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T + ) + translation -= ( + e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) + + l2e_t @ np.linalg.inv(l2e_r_mat).T + ) + sweep_cam["sensor2lidar_rotation"] = rotation.T + sweep_cam["sensor2lidar_translation"] = translation + + lidar2cam_r = np.linalg.inv(sweep_cam["sensor2lidar_rotation"]) + lidar2cam_t = sweep_cam["sensor2lidar_translation"] @ lidar2cam_r.T + lidar2cam_rt = np.eye(4) + lidar2cam_rt[:3, :3] = lidar2cam_r.T + lidar2cam_rt[3, :3] = -lidar2cam_t + intrinsic = np.array(sweep_cam["cam_intrinsic"]) + viewpad = np.eye(4) + viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic + lidar2img_rt = viewpad @ lidar2cam_rt.T + sweep_cam["intrinsics"] = viewpad.astype(np.float32) + sweep_cam["extrinsics"] = lidar2cam_rt.astype(np.float32) + sweep_cam["lidar2img"] = lidar2img_rt.astype(np.float32) + + pop_keys = [ + "ego2global_translation", + "ego2global_rotation", + "sensor2ego_translation", + "sensor2ego_rotation", + "cam_intrinsic", + ] + for key in pop_keys: + sweep_cam.pop(key) + + return sweep_cam + + +for current_id in tqdm.tqdm(range(len(key_infos["infos"]))): + e2g_t = key_infos["infos"][current_id]["ego2global_translation"] + e2g_r = key_infos["infos"][current_id]["ego2global_rotation"] + l2e_t = key_infos["infos"][current_id]["lidar2ego_translation"] + l2e_r = key_infos["infos"][current_id]["lidar2ego_rotation"] + l2e_r_mat = Quaternion(l2e_r).rotation_matrix + e2g_r_mat = Quaternion(e2g_r).rotation_matrix + + sample = nuscenes.get("sample", key_infos["infos"][current_id]["token"]) + current_cams = {} + for cam in sensors: + current_cams[cam] = nuscenes.get("sample_data", sample["data"][cam]) + + sweep_lists = [] + for _ in range(num_prev): + if sample["prev"] == "": + break + for _ in range(num_sweep): + sweep_cams = {} + for cam in sensors: + if current_cams[cam]["prev"] == "": + sweep_cams = sweep_lists[-1] + break + sample_data = nuscenes.get("sample_data", current_cams[cam]["prev"]) + sweep_cam = add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat) + current_cams[cam] = sample_data + sweep_cams[cam] = sweep_cam + sweep_lists.append(sweep_cams) + sample = nuscenes.get("sample", sample["prev"]) + sweep_cams = {} + for cam in sensors: + sample_data = nuscenes.get("sample_data", sample["data"][cam]) + sweep_cam = add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat) + current_cams[cam] = sample_data + sweep_cams[cam] = sweep_cam + sweep_lists.append(sweep_cams) + key_infos["infos"][current_id]["sweeps"] = sweep_lists + +mmcv.dump(key_infos, info_path) diff --git a/examples/onnx_ptq/petr/quantize.py b/examples/onnx_ptq/petr/quantize.py new file mode 100644 index 00000000000..80ebe8d6217 --- /dev/null +++ b/examples/onnx_ptq/petr/quantize.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import re +from pathlib import Path + +import numpy as np +import onnx +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.quantization import quantize +from modelopt.onnx.utils import topologically_sort_graph_nodes + + +class FileCalibrationReader(CalibrationDataReader): + def __init__(self, calibration_dir, onnx_path): + graph = onnx.load(onnx_path, load_external_data=False).graph + self.input_dtypes = { + value.name: onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) + for value in graph.input + } + self.batch_paths = sorted(Path(calibration_dir).glob("*.npz")) + if not self.batch_paths: + raise ValueError(f"No calibration batches found in {calibration_dir}") + self.rewind() + + def get_next(self): + batch_path = next(self._iterator, None) + return None if batch_path is None else self.load(batch_path) + + def get_first(self): + return self.load(self.batch_paths[0]) + + def rewind(self): + self._iterator = iter(self.batch_paths) + + def load(self, batch_path): + with np.load(batch_path) as batch: + missing = self.input_dtypes.keys() - batch.files + if missing: + raise ValueError(f"{batch_path} is missing inputs: {sorted(missing)}") + return { + name: batch[name].astype(dtype, copy=False) + for name, dtype in self.input_dtypes.items() + } + + +def find_backbone_nodes_to_exclude(onnx_path): + graph = onnx.load(onnx_path, load_external_data=False).graph + topologically_sort_graph_nodes(graph) + excluded = set() + downstream_tensors = set() + for node in graph.node: + is_osa = "OSA4_5" in node.name + is_downstream = any(name in downstream_tensors for name in node.input) + if is_osa or is_downstream: + excluded.add(node.name) + if "lateral_convs" in node.name or (is_downstream and not is_osa): + downstream_tensors.update(node.output) + return [rf"^{re.escape(name)}$" for name in sorted(excluded)] + + +def default_output(onnx_path, precision): + path = Path(onnx_path) + return str(path.with_name(f"{path.stem}.{precision}{path.suffix}")) + + +def quantize_model(onnx_path, calibration_dir, precision, output_path, nodes_to_exclude=()): + quantize( + onnx_path=onnx_path, + quantize_mode=precision, + calibration_data_reader=FileCalibrationReader(calibration_dir, onnx_path), + calibration_method="max", + calibration_eps=["cuda:0", "cpu"], + nodes_to_exclude=list(nodes_to_exclude), + high_precision_dtype="fp16", + output_path=output_path, + ) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Quantize PETR ONNX backbone and head models") + parser.add_argument("--backbone-onnx", required=True) + parser.add_argument("--head-onnx", required=True) + parser.add_argument("--calibration-dir", required=True, type=Path) + parser.add_argument("--precision", choices=("int8", "fp8"), default="int8") + parser.add_argument("--quantize-head", action="store_true") + parser.add_argument("--backbone-output") + parser.add_argument("--head-output") + return parser.parse_args() + + +def main(): + args = parse_args() + backbone_output = args.backbone_output or default_output(args.backbone_onnx, args.precision) + excluded = find_backbone_nodes_to_exclude(args.backbone_onnx) + print(f"Excluding {len(excluded)} accuracy-sensitive backbone nodes") + quantize_model( + args.backbone_onnx, + args.calibration_dir / "backbone", + args.precision, + backbone_output, + excluded, + ) + if args.quantize_head: + head_output = args.head_output or default_output(args.head_onnx, args.precision) + quantize_model( + args.head_onnx, + args.calibration_dir / "head", + args.precision, + head_output, + ) + else: + print("Keeping the head in FP16; use --quantize-head to quantize it") + + +if __name__ == "__main__": + main() diff --git a/examples/onnx_ptq/petr/requirements-mmdet3d.txt b/examples/onnx_ptq/petr/requirements-mmdet3d.txt new file mode 100644 index 00000000000..0b5a8fea604 --- /dev/null +++ b/examples/onnx_ptq/petr/requirements-mmdet3d.txt @@ -0,0 +1 @@ +mmdet3d==1.0.0rc6 diff --git a/examples/onnx_ptq/petr/requirements-torch.txt b/examples/onnx_ptq/petr/requirements-torch.txt new file mode 100644 index 00000000000..a66c67c31ce --- /dev/null +++ b/examples/onnx_ptq/petr/requirements-torch.txt @@ -0,0 +1,4 @@ +--extra-index-url https://download.pytorch.org/whl/cu117 + +torch==1.13.1+cu117 +torchvision==0.14.1+cu117 diff --git a/examples/onnx_ptq/petr/requirements.txt b/examples/onnx_ptq/petr/requirements.txt new file mode 100644 index 00000000000..af522454c87 --- /dev/null +++ b/examples/onnx_ptq/petr/requirements.txt @@ -0,0 +1,25 @@ +--extra-index-url https://pypi.nvidia.com +--find-links https://download.openmmlab.com/mmcv/dist/cu117/torch1.13.0/index.html + +einops +ipython<9 +lyft-dataset-sdk +mmcv-full==1.7.0 +mmdet==2.28.2 +mmsegmentation==0.30.0 +networkx>=2.2 +numba==0.57.1 +numpy<1.24 +nuscenes-devkit +onnx +onnx-graphsurgeon==0.6.1 +onnxruntime +onnxsim +opencv-python==4.5.5.64 +plyfile +pycuda +scikit-image +setuptools<81 +tensorrt-cu13-bindings==11.1.0.106 +trimesh +yapf==0.32.0 From 4ccd1573d7d22234330a97d21c3b3cd08e7cb758 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:11:39 +0000 Subject: [PATCH 2/3] Publish PETR mixed-precision results Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/onnx_ptq/petr/README.md | 15 +++------------ examples/onnx_ptq/petr/prepare_sweep_metadata.py | 8 ++------ 2 files changed, 5 insertions(+), 18 deletions(-) diff --git a/examples/onnx_ptq/petr/README.md b/examples/onnx_ptq/petr/README.md index 3575f4b844c..51acbab375a 100644 --- a/examples/onnx_ptq/petr/README.md +++ b/examples/onnx_ptq/petr/README.md @@ -179,7 +179,7 @@ The backbone quantizer preserves the accuracy-sensitive final VoVNet stage and F ## 5. Build and evaluate TensorRT engines -Build quantized graphs as strongly typed engines: +Build the quantized backbone as a strongly typed engine: ```bash precision=int8 @@ -189,13 +189,10 @@ trtexec \ --saveEngine="engines/${model}.backbone.${precision}.engine" \ --stronglyTyped \ --skipInference -trtexec \ - --onnx="onnx_files/sim_${model}.pts_bbox_head.forward.${precision}.onnx" \ - --saveEngine="engines/${model}.head.${precision}.engine" \ - --stronglyTyped \ - --skipInference ``` +If `--quantize-head` was used, build the generated head graph with the same `trtexec` options. Otherwise, use the FP16 head engine from step 3. + Evaluate any backbone/head pairing. For example, INT8 backbone with FP16 head: ```bash @@ -227,14 +224,8 @@ Results below use TensorRT 11.1.0.106 on an NVIDIA RTX 6000 Ada Generation GPU. | PETRv1-vov-p4-800x320 | FP16 | FP16 | TensorRT 11.1 | 14.507 | 0.3781 | | PETRv1-vov-p4-800x320 | INT8 | FP16 | TensorRT 11.1 | 9.992 | 0.3711 | | PETRv1-vov-p4-800x320 | FP8 | FP16 | TensorRT 11.1 | 11.455 | 0.3757 | -| PETRv1-vov-p4-800x320 | INT8 | INT8 | TensorRT 11.1 | 5.877 | 0.1779 | -| PETRv1-vov-p4-800x320 | FP8 | FP8 | TensorRT 11.1 | 8.749 | 0.2071 | | PETRv2-vov-p4-800x320 | FP16 | FP16 | TensorRT 11.1 | 19.349 | 0.4105 | | PETRv2-vov-p4-800x320 | INT8 | FP16 | TensorRT 11.1 | 14.468 | 0.4017 | | PETRv2-vov-p4-800x320 | FP8 | FP16 | TensorRT 11.1 | 16.242 | 0.4092 | -| PETRv2-vov-p4-800x320 | INT8 | INT8 | TensorRT 11.1 | 6.724 | 0.2093 | -| PETRv2-vov-p4-800x320 | FP8 | FP8 | TensorRT 11.1 | 12.261 | 0.2568 | TensorRT engines are specific to the TensorRT version and GPU architecture used to build them. These x86 results are not directly comparable with the DRIVE Orin measurements in the DL4AGX reference. - -Quantizing only the backbone retains substantially more accuracy than quantizing the detection head. Fully quantized head results are included to show the measured tradeoff rather than as recommended configurations. diff --git a/examples/onnx_ptq/petr/prepare_sweep_metadata.py b/examples/onnx_ptq/petr/prepare_sweep_metadata.py index 8a01fcd36c4..d4eb9b8905e 100644 --- a/examples/onnx_ptq/petr/prepare_sweep_metadata.py +++ b/examples/onnx_ptq/petr/prepare_sweep_metadata.py @@ -52,10 +52,7 @@ def parse_args(): info_path = args.output or os.path.join(data_root, f"mmdet3d_nuscenes_30f_infos_{info_prefix}.pkl") key_infos = pickle.load(open(os.path.join(data_root, f"nuscenes_infos_{info_prefix}.pkl"), "rb")) -if info_prefix == "test": - nuscenes_version = "v1.0-test" -else: - nuscenes_version = "v1.0-trainval" +nuscenes_version = "v1.0-test" if info_prefix == "test" else "v1.0-trainval" nuscenes = NuScenes(nuscenes_version, data_root) @@ -105,10 +102,9 @@ def add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat): intrinsic = np.array(sweep_cam["cam_intrinsic"]) viewpad = np.eye(4) viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic - lidar2img_rt = viewpad @ lidar2cam_rt.T sweep_cam["intrinsics"] = viewpad.astype(np.float32) sweep_cam["extrinsics"] = lidar2cam_rt.astype(np.float32) - sweep_cam["lidar2img"] = lidar2img_rt.astype(np.float32) + sweep_cam["lidar2img"] = (viewpad @ lidar2cam_rt.T).astype(np.float32) pop_keys = [ "ego2global_translation", From 30adf0534872bc8d79b7605b48e0191ee19b1d69 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:23:14 +0000 Subject: [PATCH 3/3] Address PETR example review feedback Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .pre-commit-config.yaml | 3 + LICENSE | 1 + examples/onnx_ptq/far3d/evaluate.py | 128 +------------- examples/onnx_ptq/far3d/quantize.py | 73 +------- examples/onnx_ptq/petr/Dockerfile | 18 +- examples/onnx_ptq/petr/README.md | 96 +++++++---- examples/onnx_ptq/petr/convert_to_fp16.py | 43 ----- examples/onnx_ptq/petr/evaluate.py | 134 ++++----------- examples/onnx_ptq/petr/prepare_calibration.py | 21 ++- .../onnx_ptq/petr/prepare_sweep_metadata.py | 157 ++++++++++-------- examples/onnx_ptq/petr/quantize.py | 60 +------ examples/onnx_ptq/petr/requirements.txt | 32 ++-- examples/onnx_ptq/quantization_utils.py | 84 ++++++++++ examples/onnx_ptq/trt_runner.py | 148 +++++++++++++++++ 14 files changed, 480 insertions(+), 518 deletions(-) delete mode 100644 examples/onnx_ptq/petr/convert_to_fp16.py create mode 100644 examples/onnx_ptq/quantization_utils.py create mode 100644 examples/onnx_ptq/trt_runner.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85b5577494c..b78f3782b62 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -127,6 +127,9 @@ repos: examples/llm_eval/mmlu.py| examples/llm_eval/modeling.py| examples/onnx_ptq/far3d/evaluate.py| + examples/onnx_ptq/petr/evaluate.py| + examples/onnx_ptq/petr/prepare_sweep_metadata.py| + examples/onnx_ptq/trt_runner.py| examples/llm_qat/train.py| examples/llm_sparsity/weight_sparsity/finetune.py| examples/specdec_bench/specdec_bench/models/specbench_medusa.py| diff --git a/LICENSE b/LICENSE index a894d488493..c58bddda878 100644 --- a/LICENSE +++ b/LICENSE @@ -224,6 +224,7 @@ the following copyright holders, licensed under the Apache License, Version 2.0 Copyright (c) 2024 Heming Xia Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team Copyright (c) OpenMMLab. All rights reserved. + Copyright (c) 2022 megvii-model. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use these files except in compliance with the License. You may obtain a copy diff --git a/examples/onnx_ptq/far3d/evaluate.py b/examples/onnx_ptq/far3d/evaluate.py index 3fac0a08013..a858d711fea 100644 --- a/examples/onnx_ptq/far3d/evaluate.py +++ b/examples/onnx_ptq/far3d/evaluate.py @@ -21,9 +21,10 @@ import argparse import importlib import os +import sys import warnings +from pathlib import Path -import tensorrt as trt import torch from mmcv import Config, DictAction from mmcv.utils import import_modules_from_strings @@ -33,115 +34,9 @@ from projects.mmdet3d_plugin.datasets.builder import build_dataloader from tqdm import tqdm -TRT_TO_TORCH = { - trt.DataType.FLOAT: torch.float32, - trt.DataType.HALF: torch.float16, - trt.DataType.INT8: torch.int8, - trt.DataType.INT32: torch.int32, - trt.DataType.BOOL: torch.bool, - trt.DataType.UINT8: torch.uint8, -} -if int(trt.__version__.split(".")[0]) >= 10: - TRT_TO_TORCH[trt.DataType.INT64] = torch.int64 - -TRT_LOGGER = trt.Logger(trt.Logger.WARNING) -trt.init_libnvinfer_plugins(TRT_LOGGER, "") - - -def aligned_tensor(shape, dtype, device, alignment=256): - element_size = torch.empty((), dtype=dtype).element_size() - element_count = int(torch.tensor(shape).prod().item()) - storage = torch.empty(element_count + alignment // element_size, dtype=dtype, device=device) - offset_bytes = (-storage.data_ptr()) % alignment - offset = offset_bytes // element_size - return storage[offset : offset + element_count].view(shape) - - -class TensorRTRunner: - def __init__(self, engine_path, state_names=()): - with open(engine_path, "rb") as engine_file: - engine_bytes = engine_file.read() - self.engine = trt.Runtime(TRT_LOGGER).deserialize_cuda_engine(engine_bytes) - if self.engine is None: - raise RuntimeError(f"Failed to deserialize {engine_path}") - self.context = self.engine.create_execution_context() - if self.context is None: - raise RuntimeError(f"Failed to create an execution context for {engine_path}") - self.tensor_names = [ - self.engine.get_tensor_name(index) for index in range(self.engine.num_io_tensors) - ] - self.input_shapes = {} - self.output_shapes = {} - self.tensor_dtypes = {} - for name in self.tensor_names: - shape = tuple(self.engine.get_tensor_shape(name)) - dtype = TRT_TO_TORCH[self.engine.get_tensor_dtype(name)] - self.tensor_dtypes[name] = dtype - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: - self.input_shapes[name] = shape - else: - self.output_shapes[name] = shape - - self.state = {} - for base_name in state_names: - name = self.resolve_name(base_name) - if name in self.input_shapes: - tensor = aligned_tensor(self.input_shapes[name], self.tensor_dtypes[name], "cuda") - tensor.zero_() - self.state[name] = tensor - self.context.set_tensor_address(name, tensor.data_ptr()) - if self.state: - torch.cuda.synchronize() - - def resolve_name(self, base_name): - if base_name in self.tensor_names: - return base_name - suffixed_name = f"{base_name}.1" - return suffixed_name if suffixed_name in self.tensor_names else base_name - - def reset_state(self): - for tensor in self.state.values(): - tensor.zero_() - - def prepare_input(self, name, inputs): - shape = self.input_shapes[name] - base_name = name.rsplit(".1", maxsplit=1)[0] if name.endswith(".1") else name - if base_name not in inputs: - raise KeyError(f"Missing TensorRT input {base_name}") - value = inputs[base_name].to(device="cuda", dtype=self.tensor_dtypes[name]) - if tuple(value.shape) != shape: - if tuple(value.shape[1:]) == shape: - value = value.squeeze(0) - elif tuple(shape[1:]) == tuple(value.shape): - value = value.unsqueeze(0) - else: - raise ValueError( - f"Input {base_name} has shape {tuple(value.shape)}, expected {shape}" - ) - return value - - def __call__(self, stream, **inputs): - input_buffers = {} - for name, shape in self.input_shapes.items(): - if name in self.state: - continue - value = self.prepare_input(name, inputs) - buffer = aligned_tensor(shape, value.dtype, value.device) - buffer.copy_(value) - input_buffers[name] = buffer - self.context.set_tensor_address(name, buffer.data_ptr()) - - outputs = {} - for name, shape in self.output_shapes.items(): - output = aligned_tensor(shape, self.tensor_dtypes[name], "cuda") - outputs[name] = output - self.context.set_tensor_address(name, output.data_ptr()) - - if not self.context.execute_async_v3(stream.cuda_stream): - raise RuntimeError("TensorRT execution failed") - stream.synchronize() - return outputs +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) +from examples.onnx_ptq.trt_runner import TensorRTRunner STATE_NAMES = ( "memory_embedding", @@ -154,8 +49,7 @@ def __call__(self, stream, **inputs): class Far3DDecoderRunner(TensorRTRunner): def __init__(self, engine_path, input_callback=None): - super().__init__(engine_path, STATE_NAMES) - self.input_callback = input_callback + super().__init__(engine_path, STATE_NAMES, input_callback) self.scene_token = None self.timestamp_offset = None @@ -175,16 +69,6 @@ def __call__(self, stream, img_metas, timestamp, **inputs): device="cuda", ) inputs["timestamp"] = (timestamp - self.timestamp_offset).float() - if self.input_callback: - calibration_inputs = {} - for name in self.input_shapes: - base_name = name.rsplit(".1", maxsplit=1)[0] if name.endswith(".1") else name - if name in self.state: - value = self.state[name] - else: - value = self.prepare_input(name, inputs) - calibration_inputs[base_name] = value - self.input_callback(calibration_inputs) outputs = super().__call__(stream, **inputs) for base_name in STATE_NAMES: input_name = self.resolve_name(base_name) @@ -287,7 +171,7 @@ def main(): } } ) - if args.max_samples is not None and len(outputs) == args.max_samples: + if args.max_samples is not None and len(outputs) >= args.max_samples: break if len(outputs) < len(dataset): diff --git a/examples/onnx_ptq/far3d/quantize.py b/examples/onnx_ptq/far3d/quantize.py index 3b844d6988f..ae7b9ba832f 100644 --- a/examples/onnx_ptq/far3d/quantize.py +++ b/examples/onnx_ptq/far3d/quantize.py @@ -14,36 +14,19 @@ # limitations under the License. import argparse -import re +import sys from pathlib import Path import numpy as np -import onnx -from onnxruntime.quantization.calibrate import CalibrationDataReader from modelopt.onnx.quantization import quantize -from modelopt.onnx.utils import topologically_sort_graph_nodes - -class FileCalibrationReader(CalibrationDataReader): - def __init__(self, calibration_dir, pattern): - self.batch_paths = sorted(Path(calibration_dir).glob(pattern)) - if not self.batch_paths: - raise ValueError(f"No {pattern} calibration batches found in {calibration_dir}") - self.rewind() - - def get_next(self): - batch_path = next(self._iterator, None) - return None if batch_path is None else self.load(batch_path) - - def get_first(self): - return self.load(self.batch_paths[0]) - - def rewind(self): - self._iterator = iter(self.batch_paths) - - def load(self, batch_path): - raise NotImplementedError +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from quantization_utils import ( + FileCalibrationReader, + NpzCalibrationReader, + find_vovnet_nodes_to_exclude, +) class EncoderCalibrationReader(FileCalibrationReader): @@ -54,42 +37,6 @@ def load(self, batch_path): return {"img": np.load(batch_path)} -class DecoderCalibrationReader(FileCalibrationReader): - def __init__(self, calibration_dir, onnx_path): - graph = onnx.load(onnx_path, load_external_data=False).graph - self.input_dtypes = { - value.name: onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) - for value in graph.input - } - super().__init__(calibration_dir, "*.npz") - - def load(self, batch_path): - with np.load(batch_path) as batch: - missing = self.input_dtypes.keys() - batch.files - if missing: - raise ValueError(f"{batch_path} is missing decoder inputs: {sorted(missing)}") - return { - name: batch[name].astype(dtype, copy=False) - for name, dtype in self.input_dtypes.items() - } - - -def find_encoder_nodes_to_exclude(onnx_path): - graph = onnx.load(onnx_path, load_external_data=False).graph - topologically_sort_graph_nodes(graph) - - excluded = set() - downstream_tensors = set() - for node in graph.node: - is_osa = "OSA4_5" in node.name - is_downstream = any(name in downstream_tensors for name in node.input) - if is_osa or is_downstream: - excluded.add(node.name) - if "lateral_convs" in node.name or (is_downstream and not is_osa): - downstream_tensors.update(node.output) - return sorted(excluded) - - def parse_args(): parser = argparse.ArgumentParser(description="Quantize the FAR3D ONNX models") parser.add_argument("--encoder-onnx", required=True, help="Path to far3d.encoder.onnx") @@ -112,9 +59,7 @@ def quantize_encoder(args): encoder_dir = Path(args.calibration_dir) if (encoder_dir / "encoder").is_dir(): encoder_dir /= "encoder" - excluded_nodes = [ - rf"^{re.escape(name)}$" for name in find_encoder_nodes_to_exclude(args.encoder_onnx) - ] + excluded_nodes = find_vovnet_nodes_to_exclude(args.encoder_onnx) print(f"Excluding {len(excluded_nodes)} accuracy-sensitive nodes from quantization") quantize( onnx_path=args.encoder_onnx, @@ -133,7 +78,7 @@ def quantize_decoder(args): quantize( onnx_path=args.decoder_onnx, quantize_mode=args.quantization_mode, - calibration_data_reader=DecoderCalibrationReader(decoder_dir, args.decoder_onnx), + calibration_data_reader=NpzCalibrationReader(decoder_dir, args.decoder_onnx), calibration_method="max", calibration_eps=["cuda:0", "cpu"], high_precision_dtype="fp16" if args.quantization_mode == "fp8" else "fp32", diff --git a/examples/onnx_ptq/petr/Dockerfile b/examples/onnx_ptq/petr/Dockerfile index 37333a86f7c..f7a394bf5a1 100644 --- a/examples/onnx_ptq/petr/Dockerfile +++ b/examples/onnx_ptq/petr/Dockerfile @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. FROM nvcr.io/nvidia/pytorch:26.07-py3 @@ -12,9 +24,9 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins ENV UV_PYTHON_INSTALL_DIR=/opt/python -RUN python -m pip install --no-cache-dir uv && \ - uv python install 3.8 && \ - uv venv --seed --python 3.8 /opt/petr +RUN python -m pip install --no-cache-dir uv==0.12.3 && \ + uv python install 3.8.20 && \ + uv venv --seed --python 3.8.20 /opt/petr COPY examples/onnx_ptq/petr/requirements*.txt /tmp/petr-requirements/ RUN env -u PIP_CONSTRAINT /opt/petr/bin/python -m pip install --no-cache-dir \ diff --git a/examples/onnx_ptq/petr/README.md b/examples/onnx_ptq/petr/README.md index 51acbab375a..a6aa0717237 100644 --- a/examples/onnx_ptq/petr/README.md +++ b/examples/onnx_ptq/petr/README.md @@ -10,16 +10,17 @@ Clone the reference repositories and apply the DL4AGX compatibility patch: ```bash git clone https://github.com/NVIDIA/DL4AGX.git +git -C DL4AGX checkout 9f7b29104c253d5bc68334e7b83b3eecb72d4572 git clone https://github.com/megvii-research/PETR.git -cd PETR -git apply ../DL4AGX/AV-Solutions/petr-trt/patch.diff -git clone https://github.com/open-mmlab/mmdetection3d.git -b v0.17.1 -mkdir -p ckpts data -ln -s /path/to/nuscenes data/nuscenes -cd .. +git -C PETR checkout f7525f93467a33707ef401c587a52d5e7b34de74 +git -C PETR apply ../DL4AGX/AV-Solutions/petr-trt/patch.diff +git clone https://github.com/open-mmlab/mmdetection3d.git PETR/mmdetection3d +git -C PETR/mmdetection3d checkout f1107977dfd26155fc1f83779ee6535d2468f449 +mkdir -p PETR/ckpts PETR/data +ln -s /data/Dataset/nuScenes PETR/data/nuscenes ``` -Download the `PETR-vov-p4-800x320_epoch24.pth` and `PETRv2-vov-p4-800x320_epoch24.pth` checkpoints linked from the DL4AGX README, rename them as shown below, and prepare this layout: +Download the `PETR-vov-p4-800x320_epoch24.pth` and `PETRv2-vov-p4-800x320_epoch24.pth` checkpoints linked from the DL4AGX README and rename them as shown below. Use nuScenes only under its [terms of use](https://www.nuscenes.org/terms-of-use). ```text PETR/ @@ -27,24 +28,13 @@ PETR/ │ ├── PETR-vov-p4-800x320_e24.pth │ └── PETRv2-vov-p4-800x320_e24.pth ├── data/nuscenes/ +│ ├── maps/ │ ├── samples/ │ ├── sweeps/ -│ ├── v1.0-trainval/ -│ └── nuscenes_infos_val.pkl +│ └── v1.0-trainval/ └── mmdetection3d/ ``` -PETRv2 also needs metadata for the previous camera sweeps: - -```bash -/opt/petr/bin/python \ - /opt/Model-Optimizer/examples/onnx_ptq/petr/prepare_sweep_metadata.py \ - /workspace/PETR/data/nuscenes \ - --split val -``` - -This creates `mmdet3d_nuscenes_30f_infos_val.pkl`. - Build and start the example image from the Model Optimizer checkout: ```bash @@ -53,10 +43,13 @@ docker build \ -t petr-modelopt \ . -docker run --rm -it --network=host --gpus=all --shm-size=64G \ +docker run --rm -it --gpus=all --shm-size=64G \ + --user "$(id -u):$(id -g)" \ + -e HOME=/tmp \ -v /path/to/Model-Optimizer:/opt/Model-Optimizer \ -v /path/to/PETR:/workspace/PETR \ -v /path/to/DL4AGX:/workspace/DL4AGX \ + -v /path/to/nuscenes:/data/Dataset/nuScenes \ petr-modelopt ``` @@ -66,6 +59,33 @@ Use the isolated Python 3.8 environment for PETR export, calibration preparation export PYTHONPATH=/opt/petr/lib/python3.8/site-packages:/workspace/PETR:/workspace/DL4AGX/AV-Solutions/petr-trt/export_eval ``` +Inside the container, generate `nuscenes_infos_val.pkl` from the local nuScenes tables with the converter included in the pinned PETR checkout: + +```bash +cd /workspace/PETR +/opt/petr/bin/python - <<'PY' +from tools.data_converter.nuscenes_converter import create_nuscenes_infos + +create_nuscenes_infos( + "/workspace/PETR/data/nuscenes", + "nuscenes", + version="v1.0-trainval", + max_sweeps=10, +) +PY +``` + +PETRv2 also needs metadata for the previous camera sweeps. Run this step in the same container: + +```bash +/opt/petr/bin/python \ + /opt/Model-Optimizer/examples/onnx_ptq/petr/prepare_sweep_metadata.py \ + /workspace/PETR/data/nuscenes \ + --split val +``` + +This creates `mmdet3d_nuscenes_30f_infos_val.pkl` without overwriting an existing file. + ## 2. Export PETRv1 and PETRv2 to ONNX Follow the DL4AGX exporter setup: @@ -91,10 +111,10 @@ The exporters create a backbone graph and a head graph. Simplify both graphs: ```bash for version in v1 v2; do model="PETR${version}" - python -m onnxsim \ + /opt/petr/bin/python -m onnxsim \ "onnx_files/${model}.extract_feat.onnx" \ "onnx_files/sim_${model}.extract_feat.onnx" - python -m onnxsim \ + /opt/petr/bin/python -m onnxsim \ "onnx_files/${model}.pts_bbox_head.forward.onnx" \ "onnx_files/sim_${model}.pts_bbox_head.forward.onnx" done @@ -135,21 +155,27 @@ Collect 512 representative backbone and head input batches. The default interval Repeat with `v2`, the PETRv2 config, checkpoint, graphs, engines, and `calibration/PETRv2` output directory. TensorRT 11 uses strongly typed networks and no longer accepts `--fp16`. -Create calibrated mixed-FP16 graphs with Model Optimizer AutoCast, preserving -FP32 model inputs and outputs, then build the reference engines. The example -below uses one representative batch for AutoCast node classification: +Use the base Python environment for Model Optimizer AutoCast, preserving FP32 +model inputs and outputs, then build the reference engines. The example below +uses one representative batch for AutoCast node classification: ```bash for version in v1 v2; do model="PETR${version}" - python /opt/Model-Optimizer/examples/onnx_ptq/petr/convert_to_fp16.py \ - "onnx_files/sim_${model}.extract_feat.onnx" \ - "onnx_files/sim_${model}.extract_feat.fp16.onnx" \ - "calibration/${model}/backbone/batch_0000.npz" - python /opt/Model-Optimizer/examples/onnx_ptq/petr/convert_to_fp16.py \ - "onnx_files/sim_${model}.pts_bbox_head.forward.onnx" \ - "onnx_files/sim_${model}.pts_bbox_head.forward.fp16.onnx" \ - "calibration/${model}/head/batch_0000.npz" + env -u PYTHONPATH python -m modelopt.onnx.autocast \ + --onnx_path "onnx_files/sim_${model}.extract_feat.onnx" \ + --output_path "onnx_files/sim_${model}.extract_feat.fp16.onnx" \ + --calibration_data "calibration/${model}/backbone/batch_0000.npz" \ + --low_precision_type fp16 \ + --keep_io_types \ + --providers cuda:0 cpu + env -u PYTHONPATH python -m modelopt.onnx.autocast \ + --onnx_path "onnx_files/sim_${model}.pts_bbox_head.forward.onnx" \ + --output_path "onnx_files/sim_${model}.pts_bbox_head.forward.fp16.onnx" \ + --calibration_data "calibration/${model}/head/batch_0000.npz" \ + --low_precision_type fp16 \ + --keep_io_types \ + --providers cuda:0 cpu trtexec \ --onnx="onnx_files/sim_${model}.extract_feat.fp16.onnx" \ --saveEngine="engines/${model}.backbone.fp16.engine" \ @@ -166,7 +192,7 @@ done Use the base Python environment for Model Optimizer. This command quantizes the backbone and keeps the head in FP16: ```bash -python /opt/Model-Optimizer/examples/onnx_ptq/petr/quantize.py \ +env -u PYTHONPATH python /opt/Model-Optimizer/examples/onnx_ptq/petr/quantize.py \ --backbone-onnx onnx_files/sim_PETRv1.extract_feat.onnx \ --head-onnx onnx_files/sim_PETRv1.pts_bbox_head.forward.onnx \ --calibration-dir calibration/PETRv1 \ diff --git a/examples/onnx_ptq/petr/convert_to_fp16.py b/examples/onnx_ptq/petr/convert_to_fp16.py deleted file mode 100644 index d2fc0c5aa88..00000000000 --- a/examples/onnx_ptq/petr/convert_to_fp16.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import argparse - -from modelopt.onnx import utils as onnx_utils -from modelopt.onnx.autocast import convert_to_mixed_precision - - -def parse_args(): - parser = argparse.ArgumentParser(description="Convert a PETR ONNX model to FP16") - parser.add_argument("input") - parser.add_argument("output") - parser.add_argument("calibration_data") - return parser.parse_args() - - -def main(): - args = parse_args() - model = convert_to_mixed_precision( - onnx_path=args.input, - low_precision_type="fp16", - calibration_data=args.calibration_data, - keep_io_types=True, - providers=["cuda:0", "cpu"], - ) - onnx_utils.save_onnx(model, args.output) - - -if __name__ == "__main__": - main() diff --git a/examples/onnx_ptq/petr/evaluate.py b/examples/onnx_ptq/petr/evaluate.py index 31ed8926ecd..c50c13e7264 100644 --- a/examples/onnx_ptq/petr/evaluate.py +++ b/examples/onnx_ptq/petr/evaluate.py @@ -1,4 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from https://github.com/NVIDIA/DL4AGX/blob/9f7b29104c253d5bc68334e7b83b3eecb72d4572/AV-Solutions/petr-trt/export_eval/v1/v1_evaluate_trt.py +# and https://github.com/NVIDIA/DL4AGX/blob/9f7b29104c253d5bc68334e7b83b3eecb72d4572/AV-Solutions/petr-trt/export_eval/v2/v2_evaluate_trt.py. +# Copyright (c) OpenMMLab. All rights reserved. +# +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,14 +17,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Adapted from https://github.com/NVIDIA/DL4AGX/tree/9f7b29104c253d5bc68334e7b83b3eecb72d4572/AV-Solutions/petr-trt/export_eval. -# Copyright (c) OpenMMLab. All rights reserved. - import argparse import importlib import os +import sys +from pathlib import Path -import tensorrt as trt import torch import torch.nn.functional as F from mmcv import Config, DictAction @@ -32,94 +34,17 @@ from mmdet3d.models import build_model from tqdm import tqdm -TRT_TO_TORCH = { - trt.DataType.FLOAT: torch.float32, - trt.DataType.HALF: torch.float16, - trt.DataType.INT8: torch.int8, - trt.DataType.INT32: torch.int32, - trt.DataType.BOOL: torch.bool, - trt.DataType.UINT8: torch.uint8, -} -if int(trt.__version__.split(".")[0]) >= 10: - TRT_TO_TORCH[trt.DataType.INT64] = torch.int64 - -TRT_LOGGER = trt.Logger(trt.Logger.WARNING) -trt.init_libnvinfer_plugins(TRT_LOGGER, "") - - -def aligned_tensor(shape, dtype, device, alignment=256): - element_size = torch.empty((), dtype=dtype).element_size() - element_count = int(torch.tensor(shape).prod().item()) - storage = torch.empty(element_count + alignment // element_size, dtype=dtype, device=device) - offset = ((-storage.data_ptr()) % alignment) // element_size - return storage[offset : offset + element_count].view(shape) - - -class TensorRTRunner: - def __init__(self, engine_path, input_callback=None): - with open(engine_path, "rb") as engine_file: - engine_bytes = engine_file.read() - self.engine = trt.Runtime(TRT_LOGGER).deserialize_cuda_engine(engine_bytes) - if self.engine is None: - raise RuntimeError(f"Failed to deserialize {engine_path}") - self.context = self.engine.create_execution_context() - self.input_callback = input_callback - self.input_names = [] - self.output_names = [] - self.shapes = {} - self.dtypes = {} - for index in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(index) - self.shapes[name] = tuple(self.engine.get_tensor_shape(name)) - self.dtypes[name] = TRT_TO_TORCH[self.engine.get_tensor_dtype(name)] - names = ( - self.input_names - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT - else self.output_names - ) - names.append(name) +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) - @staticmethod - def fit_shape(value, shape): - if tuple(value.shape) == shape: - return value - if tuple(value.shape[1:]) == shape: - return value.squeeze(0) - if tuple(shape[1:]) == tuple(value.shape): - return value.unsqueeze(0) - raise ValueError(f"Input has shape {tuple(value.shape)}, expected {shape}") - - def __call__(self, stream, values): - if len(values) != len(self.input_names): - raise ValueError(f"Received {len(values)} inputs, expected {len(self.input_names)}") - inputs = [] - for name, value in zip(self.input_names, values): - value = value.to(device="cuda", dtype=self.dtypes[name]).contiguous() - value = self.fit_shape(value, self.shapes[name]) - buffer = aligned_tensor(self.shapes[name], value.dtype, value.device) - buffer.copy_(value) - inputs.append(buffer) - self.context.set_tensor_address(name, buffer.data_ptr()) - if self.input_callback: - self.input_callback(inputs) - - outputs = [] - for name in self.output_names: - output = aligned_tensor(self.shapes[name], self.dtypes[name], "cuda") - outputs.append(output) - self.context.set_tensor_address(name, output.data_ptr()) - if not self.context.execute_async_v3(stream.cuda_stream): - raise RuntimeError("TensorRT execution failed") - stream.synchronize() - return outputs +from examples.onnx_ptq.trt_runner import TensorRTRunner def import_plugin(cfg): if cfg.get("custom_imports"): import_modules_from_strings(**cfg.custom_imports) - if cfg.get("plugin"): - plugin_dir = os.path.dirname(cfg.plugin_dir).split("/") - importlib.import_module(".".join(plugin_dir)) + plugin_dir = cfg.get("plugin_dir") + if cfg.get("plugin") and plugin_dir: + importlib.import_module(".".join(os.path.dirname(plugin_dir).split("/"))) def build_runtime(config_path, checkpoint_path, cfg_options=None): @@ -153,8 +78,8 @@ class PETRPipeline: def __init__(self, version, model, backbone_engine, head_engine, callbacks=(None, None)): self.version = version self.model = model - self.backbone = TensorRTRunner(backbone_engine, callbacks[0]) - self.head = TensorRTRunner(head_engine, callbacks[1]) + self.backbone = TensorRTRunner(backbone_engine, input_callback=callbacks[0]) + self.head = TensorRTRunner(head_engine, input_callback=callbacks[1]) @staticmethod def masks(features, img_metas): @@ -169,35 +94,40 @@ def masks(features, img_metas): def backbone_inputs(self, images, img_metas): if self.version == "v1": - return [images] + return {"img": images} current = images[:, :6].contiguous() previous = images[:, 6:12].contiguous() previous_features = self.model.extract_img_feat(previous, img_metas) - return [current, *previous_features] + return { + "img": current, + **{f"prev.{index}": value for index, value in enumerate(previous_features)}, + } def head_inputs(self, features, img_metas): masks = self.masks(features, img_metas) coords, _ = self.model.pts_bbox_head.position_embeding(features, img_metas, masks) - values = [features[0]] + inputs = {"mlvl_feats.0": features[0]} if self.version == "v2": timestamps = features[0].new_tensor([meta["timestamp"] for meta in img_metas]) timestamps = timestamps.view(1, -1, 6) - values.append((timestamps[:, 1] - timestamps[:, 0]).mean(-1)) - values.append(coords) - return values + inputs["img_metas.0[mean_time_stamp]"] = (timestamps[:, 1] - timestamps[:, 0]).mean(-1) + inputs["img_metas.0[coords_position_embeding]"] = coords + return inputs def __call__(self, stream, data): images = data["img"][0].data[0].cuda() img_metas = data["img_metas"][0].data[0] with torch.cuda.stream(stream), torch.no_grad(): - features = self.backbone(stream, self.backbone_inputs(images, img_metas)) + feature_outputs = self.backbone(stream, **self.backbone_inputs(images, img_metas)) camera_count = 6 if self.version == "v1" else 12 - features = [value.reshape(1, camera_count, *value.shape[-3:]) for value in features] - outputs = self.head(stream, self.head_inputs(features, img_metas)) - outputs = [value.float() for value in outputs] + features = [ + feature_outputs[name].reshape(1, camera_count, *feature_outputs[name].shape[-3:]) + for name in ("out.0", "out.1") + ] + outputs = self.head(stream, **self.head_inputs(features, img_metas)) head_outputs = { - "all_cls_scores": outputs[0], - "all_bbox_preds": outputs[1], + "all_cls_scores": outputs["out.all_cls_scores"].float(), + "all_bbox_preds": outputs["out.all_bbox_preds"].float(), "enc_cls_scores": None, "enc_bbox_preds": None, } @@ -233,7 +163,7 @@ def main(): outputs = [] for data in tqdm(loader): outputs.extend(pipeline(stream, data)) - if args.max_samples is not None and len(outputs) == args.max_samples: + if args.max_samples is not None and len(outputs) >= args.max_samples: break if len(outputs) < len(dataset): print(f"Processed {len(outputs)} samples; skipping dataset metrics") diff --git a/examples/onnx_ptq/petr/prepare_calibration.py b/examples/onnx_ptq/petr/prepare_calibration.py index 7d8d00c9911..614307ca653 100644 --- a/examples/onnx_ptq/petr/prepare_calibration.py +++ b/examples/onnx_ptq/petr/prepare_calibration.py @@ -31,19 +31,23 @@ def __init__(self, output_dir, onnx_path): if any(self.output_dir.glob("*.npz")): raise FileExistsError(f"{self.output_dir} already contains calibration batches") graph = onnx.load(onnx_path, load_external_data=False).graph - self.names = [value.name for value in graph.input] - self.dtypes = [ - onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) + self.dtypes = { + value.name: onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) for value in graph.input - ] + } self.saved = 0 def __call__(self, values): - if len(values) != len(self.names): - raise ValueError(f"Received {len(values)} inputs, expected {len(self.names)}") + missing = self.dtypes.keys() - values.keys() + unexpected = values.keys() - self.dtypes.keys() + if missing or unexpected: + raise ValueError( + f"Calibration input mismatch; missing={sorted(missing)}, " + f"unexpected={sorted(unexpected)}" + ) batch = { - name: value.detach().cpu().numpy().astype(dtype, copy=False) - for name, dtype, value in zip(self.names, self.dtypes, values) + name: values[name].detach().cpu().numpy().astype(dtype, copy=False) + for name, dtype in self.dtypes.items() } np.savez(self.output_dir / f"batch_{self.saved:04d}.npz", **batch) self.saved += 1 @@ -101,5 +105,4 @@ def main(): if __name__ == "__main__": - torch.multiprocessing.set_start_method("fork") main() diff --git a/examples/onnx_ptq/petr/prepare_sweep_metadata.py b/examples/onnx_ptq/petr/prepare_sweep_metadata.py index d4eb9b8905e..31231b9ca09 100644 --- a/examples/onnx_ptq/petr/prepare_sweep_metadata.py +++ b/examples/onnx_ptq/petr/prepare_sweep_metadata.py @@ -1,3 +1,6 @@ +# Adapted from https://github.com/megvii-research/PETR/blob/f7525f93467a33707ef401c587a52d5e7b34de74/tools/generate_sweep_pkl.py. +# Copyright (c) 2022 megvii-model. All Rights Reserved. +# # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -13,12 +16,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Adapted from PETR tools/generate_sweep_pkl.py. -# Copyright (c) 2022 Megvii Inc. All rights reserved. - import argparse import os -import pickle import mmcv import numpy as np @@ -26,7 +25,7 @@ from nuscenes import NuScenes from pyquaternion import Quaternion -sensors = [ +SENSORS = [ "CAM_FRONT", "CAM_FRONT_RIGHT", "CAM_BACK_RIGHT", @@ -34,6 +33,8 @@ "CAM_BACK_LEFT", "CAM_FRONT_LEFT", ] +NUM_PREV = 5 +NUM_SWEEPS = 5 def parse_args(): @@ -44,25 +45,14 @@ def parse_args(): return parser.parse_args() -args = parse_args() -info_prefix = args.split -data_root = os.path.abspath(args.data_root) + os.sep -num_prev = 5 -num_sweep = 5 - -info_path = args.output or os.path.join(data_root, f"mmdet3d_nuscenes_30f_infos_{info_prefix}.pkl") -key_infos = pickle.load(open(os.path.join(data_root, f"nuscenes_infos_{info_prefix}.pkl"), "rb")) -nuscenes_version = "v1.0-test" if info_prefix == "test" else "v1.0-trainval" -nuscenes = NuScenes(nuscenes_version, data_root) - - -def add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat): - sweep_cam = {} - sweep_cam["is_key_frame"] = sample_data["is_key_frame"] - sweep_cam["data_path"] = os.path.join(data_root, sample_data["filename"]) - sweep_cam["type"] = "camera" - sweep_cam["timestamp"] = sample_data["timestamp"] - sweep_cam["sample_data_token"] = sample_data["sample_token"] +def add_frame(nuscenes, data_root, sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat): + sweep_cam = { + "is_key_frame": sample_data["is_key_frame"], + "data_path": os.path.join(data_root, sample_data["filename"]), + "type": "camera", + "timestamp": sample_data["timestamp"], + "sample_data_token": sample_data["sample_token"], + } pose_record = nuscenes.get("ego_pose", sample_data["ego_pose_token"]) calibrated_sensor_record = nuscenes.get( "calibrated_sensor", sample_data["calibrated_sensor_token"] @@ -74,13 +64,10 @@ def add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat): sweep_cam["sensor2ego_rotation"] = calibrated_sensor_record["rotation"] sweep_cam["cam_intrinsic"] = calibrated_sensor_record["camera_intrinsic"] - l2e_r_s = sweep_cam["sensor2ego_rotation"] - l2e_t_s = sweep_cam["sensor2ego_translation"] - e2g_r_s = sweep_cam["ego2global_rotation"] + l2e_r_s_mat = Quaternion(sweep_cam["sensor2ego_rotation"]).rotation_matrix + e2g_r_s_mat = Quaternion(sweep_cam["ego2global_rotation"]).rotation_matrix e2g_t_s = sweep_cam["ego2global_translation"] - - l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix - e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix + l2e_t_s = sweep_cam["sensor2ego_translation"] rotation = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T ) @@ -106,55 +93,87 @@ def add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat): sweep_cam["extrinsics"] = lidar2cam_rt.astype(np.float32) sweep_cam["lidar2img"] = (viewpad @ lidar2cam_rt.T).astype(np.float32) - pop_keys = [ + for key in ( "ego2global_translation", "ego2global_rotation", "sensor2ego_translation", "sensor2ego_rotation", "cam_intrinsic", - ] - for key in pop_keys: + ): sweep_cam.pop(key) return sweep_cam -for current_id in tqdm.tqdm(range(len(key_infos["infos"]))): - e2g_t = key_infos["infos"][current_id]["ego2global_translation"] - e2g_r = key_infos["infos"][current_id]["ego2global_rotation"] - l2e_t = key_infos["infos"][current_id]["lidar2ego_translation"] - l2e_r = key_infos["infos"][current_id]["lidar2ego_rotation"] - l2e_r_mat = Quaternion(l2e_r).rotation_matrix - e2g_r_mat = Quaternion(e2g_r).rotation_matrix - - sample = nuscenes.get("sample", key_infos["infos"][current_id]["token"]) - current_cams = {} - for cam in sensors: - current_cams[cam] = nuscenes.get("sample_data", sample["data"][cam]) - - sweep_lists = [] - for _ in range(num_prev): - if sample["prev"] == "": - break - for _ in range(num_sweep): +def add_sweeps(key_infos, nuscenes, data_root): + for current_id in tqdm.tqdm(range(len(key_infos["infos"]))): + info = key_infos["infos"][current_id] + e2g_t = info["ego2global_translation"] + l2e_t = info["lidar2ego_translation"] + l2e_r_mat = Quaternion(info["lidar2ego_rotation"]).rotation_matrix + e2g_r_mat = Quaternion(info["ego2global_rotation"]).rotation_matrix + + sample = nuscenes.get("sample", info["token"]) + current_cams = {cam: nuscenes.get("sample_data", sample["data"][cam]) for cam in SENSORS} + sweep_lists = [] + for _ in range(NUM_PREV): + if not sample["prev"]: + break + for _ in range(NUM_SWEEPS): + sweep_cams = {} + for cam in SENSORS: + if not current_cams[cam]["prev"]: + sweep_cams = sweep_lists[-1] if sweep_lists else {} + break + sample_data = nuscenes.get("sample_data", current_cams[cam]["prev"]) + sweep_cams[cam] = add_frame( + nuscenes, + data_root, + sample_data, + e2g_t, + l2e_t, + l2e_r_mat, + e2g_r_mat, + ) + current_cams[cam] = sample_data + sweep_lists.append(sweep_cams) + + sample = nuscenes.get("sample", sample["prev"]) sweep_cams = {} - for cam in sensors: - if current_cams[cam]["prev"] == "": - sweep_cams = sweep_lists[-1] - break - sample_data = nuscenes.get("sample_data", current_cams[cam]["prev"]) - sweep_cam = add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat) + for cam in SENSORS: + sample_data = nuscenes.get("sample_data", sample["data"][cam]) + sweep_cams[cam] = add_frame( + nuscenes, + data_root, + sample_data, + e2g_t, + l2e_t, + l2e_r_mat, + e2g_r_mat, + ) current_cams[cam] = sample_data - sweep_cams[cam] = sweep_cam sweep_lists.append(sweep_cams) - sample = nuscenes.get("sample", sample["prev"]) - sweep_cams = {} - for cam in sensors: - sample_data = nuscenes.get("sample_data", sample["data"][cam]) - sweep_cam = add_frame(sample_data, e2g_t, l2e_t, l2e_r_mat, e2g_r_mat) - current_cams[cam] = sample_data - sweep_cams[cam] = sweep_cam - sweep_lists.append(sweep_cams) - key_infos["infos"][current_id]["sweeps"] = sweep_lists - -mmcv.dump(key_infos, info_path) + info["sweeps"] = sweep_lists + + +def main(): + args = parse_args() + data_root = os.path.abspath(args.data_root) + os.sep + input_path = os.path.join(data_root, f"nuscenes_infos_{args.split}.pkl") + output_path = args.output or os.path.join( + data_root, f"mmdet3d_nuscenes_30f_infos_{args.split}.pkl" + ) + if os.path.exists(output_path): + raise FileExistsError(f"Refusing to overwrite {output_path}") + + # This metadata must be generated locally by the documented MMDetection3D data-prep step. + key_infos = mmcv.load(input_path) + + version = "v1.0-test" if args.split == "test" else "v1.0-trainval" + nuscenes = NuScenes(version, data_root) + add_sweeps(key_infos, nuscenes, data_root) + mmcv.dump(key_infos, output_path) + + +if __name__ == "__main__": + main() diff --git a/examples/onnx_ptq/petr/quantize.py b/examples/onnx_ptq/petr/quantize.py index 80ebe8d6217..9c4d73286a9 100644 --- a/examples/onnx_ptq/petr/quantize.py +++ b/examples/onnx_ptq/petr/quantize.py @@ -14,63 +14,13 @@ # limitations under the License. import argparse -import re +import sys from pathlib import Path -import numpy as np -import onnx -from onnxruntime.quantization.calibrate import CalibrationDataReader - from modelopt.onnx.quantization import quantize -from modelopt.onnx.utils import topologically_sort_graph_nodes - - -class FileCalibrationReader(CalibrationDataReader): - def __init__(self, calibration_dir, onnx_path): - graph = onnx.load(onnx_path, load_external_data=False).graph - self.input_dtypes = { - value.name: onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) - for value in graph.input - } - self.batch_paths = sorted(Path(calibration_dir).glob("*.npz")) - if not self.batch_paths: - raise ValueError(f"No calibration batches found in {calibration_dir}") - self.rewind() - - def get_next(self): - batch_path = next(self._iterator, None) - return None if batch_path is None else self.load(batch_path) - - def get_first(self): - return self.load(self.batch_paths[0]) - - def rewind(self): - self._iterator = iter(self.batch_paths) - - def load(self, batch_path): - with np.load(batch_path) as batch: - missing = self.input_dtypes.keys() - batch.files - if missing: - raise ValueError(f"{batch_path} is missing inputs: {sorted(missing)}") - return { - name: batch[name].astype(dtype, copy=False) - for name, dtype in self.input_dtypes.items() - } - -def find_backbone_nodes_to_exclude(onnx_path): - graph = onnx.load(onnx_path, load_external_data=False).graph - topologically_sort_graph_nodes(graph) - excluded = set() - downstream_tensors = set() - for node in graph.node: - is_osa = "OSA4_5" in node.name - is_downstream = any(name in downstream_tensors for name in node.input) - if is_osa or is_downstream: - excluded.add(node.name) - if "lateral_convs" in node.name or (is_downstream and not is_osa): - downstream_tensors.update(node.output) - return [rf"^{re.escape(name)}$" for name in sorted(excluded)] +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from quantization_utils import NpzCalibrationReader, find_vovnet_nodes_to_exclude def default_output(onnx_path, precision): @@ -82,7 +32,7 @@ def quantize_model(onnx_path, calibration_dir, precision, output_path, nodes_to_ quantize( onnx_path=onnx_path, quantize_mode=precision, - calibration_data_reader=FileCalibrationReader(calibration_dir, onnx_path), + calibration_data_reader=NpzCalibrationReader(calibration_dir, onnx_path), calibration_method="max", calibration_eps=["cuda:0", "cpu"], nodes_to_exclude=list(nodes_to_exclude), @@ -106,7 +56,7 @@ def parse_args(): def main(): args = parse_args() backbone_output = args.backbone_output or default_output(args.backbone_onnx, args.precision) - excluded = find_backbone_nodes_to_exclude(args.backbone_onnx) + excluded = find_vovnet_nodes_to_exclude(args.backbone_onnx) print(f"Excluding {len(excluded)} accuracy-sensitive backbone nodes") quantize_model( args.backbone_onnx, diff --git a/examples/onnx_ptq/petr/requirements.txt b/examples/onnx_ptq/petr/requirements.txt index af522454c87..86b893c9dee 100644 --- a/examples/onnx_ptq/petr/requirements.txt +++ b/examples/onnx_ptq/petr/requirements.txt @@ -1,25 +1,25 @@ --extra-index-url https://pypi.nvidia.com --find-links https://download.openmmlab.com/mmcv/dist/cu117/torch1.13.0/index.html -einops -ipython<9 -lyft-dataset-sdk +einops==0.8.1 +ipython==8.12.3 +lyft-dataset-sdk==0.0.8 mmcv-full==1.7.0 mmdet==2.28.2 mmsegmentation==0.30.0 -networkx>=2.2 -numba==0.57.1 -numpy<1.24 -nuscenes-devkit -onnx +networkx==2.2 +numba==0.53.0 +numpy==1.23.5 +nuscenes-devkit==1.1.11 +onnx==1.17.0 onnx-graphsurgeon==0.6.1 -onnxruntime -onnxsim -opencv-python==4.5.5.64 -plyfile -pycuda -scikit-image -setuptools<81 +onnxruntime==1.19.2 +onnxsim==0.5.0 +opencv-python==4.8.1.78 +plyfile==1.0.3 +pycuda==2026.1 +scikit-image==0.19.3 +setuptools==75.3.4 tensorrt-cu13-bindings==11.1.0.106 -trimesh +trimesh==2.35.39 yapf==0.32.0 diff --git a/examples/onnx_ptq/quantization_utils.py b/examples/onnx_ptq/quantization_utils.py new file mode 100644 index 00000000000..74597fbbcf9 --- /dev/null +++ b/examples/onnx_ptq/quantization_utils.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +from pathlib import Path + +import numpy as np +import onnx +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.utils import topologically_sort_graph_nodes + + +class FileCalibrationReader(CalibrationDataReader): + def __init__(self, calibration_dir, pattern): + self.batch_paths = sorted(Path(calibration_dir).glob(pattern)) + if not self.batch_paths: + raise ValueError(f"No {pattern} calibration batches found in {calibration_dir}") + self.rewind() + + def get_next(self): + batch_path = next(self._iterator, None) + return None if batch_path is None else self.load(batch_path) + + def get_first(self): + return self.load(self.batch_paths[0]) + + def rewind(self): + self._iterator = iter(self.batch_paths) + + def load(self, batch_path): + raise NotImplementedError + + +class NpzCalibrationReader(FileCalibrationReader): + def __init__(self, calibration_dir, onnx_path): + graph = onnx.load(onnx_path, load_external_data=False).graph + self.input_dtypes = { + value.name: onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) + for value in graph.input + } + super().__init__(calibration_dir, "*.npz") + + def load(self, batch_path): + with np.load(batch_path) as batch: + missing = self.input_dtypes.keys() - batch.files + if missing: + raise ValueError(f"{batch_path} is missing inputs: {sorted(missing)}") + return { + name: batch[name].astype(dtype, copy=False) + for name, dtype in self.input_dtypes.items() + } + + +def find_vovnet_nodes_to_exclude(onnx_path): + """Find the VoVNet OSA4_5 stage and nodes downstream of FPN lateral_convs.""" + graph = onnx.load(onnx_path, load_external_data=False).graph + topologically_sort_graph_nodes(graph) + + excluded = set() + downstream_tensors = set() + for node in graph.node: + is_osa = "OSA4_5" in node.name + is_downstream = any(name in downstream_tensors for name in node.input) + if is_osa or is_downstream: + excluded.add(node.name) + if "lateral_convs" in node.name or (is_downstream and not is_osa): + downstream_tensors.update(node.output) + + if not excluded: + raise ValueError(f"No accuracy-sensitive VoVNet nodes found in {onnx_path}") + return [rf"^{re.escape(name)}$" for name in sorted(excluded)] diff --git a/examples/onnx_ptq/trt_runner.py b/examples/onnx_ptq/trt_runner.py new file mode 100644 index 00000000000..c1523f4bb1e --- /dev/null +++ b/examples/onnx_ptq/trt_runner.py @@ -0,0 +1,148 @@ +# Adapted from https://github.com/NVIDIA/DL4AGX/blob/9f7b29104c253d5bc68334e7b83b3eecb72d4572/AV-Solutions/far3d-trt/tools/test_tensorrt.py +# which was modified from https://github.com/megvii-research/Far3D/blob/5efb9d73a246c39fac79b3cf8c20a8e059611c3f/tools/test.py. +# Copyright (c) OpenMMLab. All rights reserved. +# Modified by Zhiqi Li. +# +# SPDX-FileCopyrightText: Copyright (c) 2023-2024, 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import tensorrt as trt +import torch + +TRT_TO_TORCH = { + trt.DataType.FLOAT: torch.float32, + trt.DataType.HALF: torch.float16, + trt.DataType.INT8: torch.int8, + trt.DataType.INT32: torch.int32, + trt.DataType.BOOL: torch.bool, + trt.DataType.UINT8: torch.uint8, +} +if int(trt.__version__.split(".")[0]) >= 10: + TRT_TO_TORCH[trt.DataType.INT64] = torch.int64 + +TRT_LOGGER = trt.Logger(trt.Logger.WARNING) +trt.init_libnvinfer_plugins(TRT_LOGGER, "") + + +def aligned_tensor(shape, dtype, device, alignment=256): + element_size = torch.empty((), dtype=dtype).element_size() + element_count = math.prod(shape) + storage = torch.empty(element_count + alignment // element_size, dtype=dtype, device=device) + offset_bytes = (-storage.data_ptr()) % alignment + offset = offset_bytes // element_size + return storage[offset : offset + element_count].view(shape) + + +class TensorRTRunner: + def __init__(self, engine_path, state_names=(), input_callback=None): + with open(engine_path, "rb") as engine_file: + engine_bytes = engine_file.read() + self.engine = trt.Runtime(TRT_LOGGER).deserialize_cuda_engine(engine_bytes) + if self.engine is None: + raise RuntimeError(f"Failed to deserialize {engine_path}") + self.context = self.engine.create_execution_context() + if self.context is None: + raise RuntimeError(f"Failed to create an execution context for {engine_path}") + self.input_callback = input_callback + self.tensor_names = [ + self.engine.get_tensor_name(index) for index in range(self.engine.num_io_tensors) + ] + self.input_shapes = {} + self.output_shapes = {} + self.tensor_dtypes = {} + for name in self.tensor_names: + shape = tuple(self.engine.get_tensor_shape(name)) + dtype = TRT_TO_TORCH[self.engine.get_tensor_dtype(name)] + self.tensor_dtypes[name] = dtype + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + self.input_shapes[name] = shape + else: + self.output_shapes[name] = shape + + self.state = {} + for base_name in state_names: + name = self.resolve_name(base_name) + if name in self.input_shapes: + tensor = aligned_tensor(self.input_shapes[name], self.tensor_dtypes[name], "cuda") + tensor.zero_() + self.state[name] = tensor + self.context.set_tensor_address(name, tensor.data_ptr()) + if self.state: + torch.cuda.synchronize() + + def resolve_name(self, base_name): + if base_name in self.tensor_names: + return base_name + suffixed_name = f"{base_name}.1" + return suffixed_name if suffixed_name in self.tensor_names else base_name + + def reset_state(self): + for tensor in self.state.values(): + tensor.zero_() + + @staticmethod + def input_key(name, inputs): + if name in inputs: + return name + base_name = name.rsplit(".1", maxsplit=1)[0] if name.endswith(".1") else name + if base_name in inputs: + return base_name + raise KeyError(f"Missing TensorRT input {name}") + + def prepare_input(self, name, inputs): + input_key = self.input_key(name, inputs) + shape = self.input_shapes[name] + value = inputs[input_key].to(device="cuda", dtype=self.tensor_dtypes[name]) + if tuple(value.shape) != shape: + if tuple(value.shape[1:]) == shape: + value = value.squeeze(0) + elif tuple(shape[1:]) == tuple(value.shape): + value = value.unsqueeze(0) + else: + raise ValueError( + f"Input {input_key} has shape {tuple(value.shape)}, expected {shape}" + ) + return input_key, value + + def __call__(self, stream, **inputs): + input_buffers = {} + callback_inputs = {} + for name, shape in self.input_shapes.items(): + if name in self.state: + callback_name = name.rsplit(".1", maxsplit=1)[0] + callback_inputs[callback_name] = self.state[name] + continue + input_key, value = self.prepare_input(name, inputs) + buffer = aligned_tensor(shape, value.dtype, value.device) + buffer.copy_(value) + input_buffers[name] = buffer + callback_inputs[input_key] = buffer + self.context.set_tensor_address(name, buffer.data_ptr()) + + if self.input_callback: + self.input_callback(callback_inputs) + + outputs = {} + for name, shape in self.output_shapes.items(): + output = aligned_tensor(shape, self.tensor_dtypes[name], "cuda") + outputs[name] = output + self.context.set_tensor_address(name, output.data_ptr()) + + if not self.context.execute_async_v3(stream.cuda_stream): + raise RuntimeError("TensorRT execution failed") + stream.synchronize() + return outputs