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/CHANGELOG.rst b/CHANGELOG.rst index 964fd8483fc..19202cba316 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add an end-to-end PETRv1 and PETRv2 ONNX PTQ example with nuScenes calibration and accuracy evaluation, INT8 and FP8 backbone quantization, FP16 heads, and TensorRT engine benchmarking. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. *Megatron Framework (M-LM / M-Bridge)* 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/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/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 new file mode 100644 index 00000000000..f7a394bf5a1 --- /dev/null +++ b/examples/onnx_ptq/petr/Dockerfile @@ -0,0 +1,49 @@ +# 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 + +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==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 \ + -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..a6aa0717237 --- /dev/null +++ b/examples/onnx_ptq/petr/README.md @@ -0,0 +1,257 @@ +# 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 -C DL4AGX checkout 9f7b29104c253d5bc68334e7b83b3eecb72d4572 +git clone https://github.com/megvii-research/PETR.git +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 and rename them as shown below. Use nuScenes only under its [terms of use](https://www.nuscenes.org/terms-of-use). + +```text +PETR/ +├── ckpts/ +│ ├── PETR-vov-p4-800x320_e24.pth +│ └── PETRv2-vov-p4-800x320_e24.pth +├── data/nuscenes/ +│ ├── maps/ +│ ├── samples/ +│ ├── sweeps/ +│ └── v1.0-trainval/ +└── mmdetection3d/ +``` + +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 --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 +``` + +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 +``` + +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: + +```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}" + /opt/petr/bin/python -m onnxsim \ + "onnx_files/${model}.extract_feat.onnx" \ + "onnx_files/sim_${model}.extract_feat.onnx" + /opt/petr/bin/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`. +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}" + 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" \ + --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 +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 \ + --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 the quantized backbone as a strongly typed engine: + +```bash +precision=int8 +model=PETRv1 +trtexec \ + --onnx="onnx_files/sim_${model}.extract_feat.${precision}.onnx" \ + --saveEngine="engines/${model}.backbone.${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 +/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 | +| 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 | + +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. diff --git a/examples/onnx_ptq/petr/evaluate.py b/examples/onnx_ptq/petr/evaluate.py new file mode 100644 index 00000000000..c50c13e7264 --- /dev/null +++ b/examples/onnx_ptq/petr/evaluate.py @@ -0,0 +1,180 @@ +# 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"); +# 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 importlib +import os +import sys +from pathlib import Path + +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 + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from examples.onnx_ptq.trt_runner import TensorRTRunner + + +def import_plugin(cfg): + if cfg.get("custom_imports"): + import_modules_from_strings(**cfg.custom_imports) + 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): + 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, input_callback=callbacks[0]) + self.head = TensorRTRunner(head_engine, input_callback=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 {"img": images} + current = images[:, :6].contiguous() + previous = images[:, 6:12].contiguous() + previous_features = self.model.extract_img_feat(previous, img_metas) + 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) + 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) + 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(): + feature_outputs = self.backbone(stream, **self.backbone_inputs(images, img_metas)) + camera_count = 6 if self.version == "v1" else 12 + 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["out.all_cls_scores"].float(), + "all_bbox_preds": outputs["out.all_bbox_preds"].float(), + "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..614307ca653 --- /dev/null +++ b/examples/onnx_ptq/petr/prepare_calibration.py @@ -0,0 +1,108 @@ +# 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.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): + 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: 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 + + +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__": + 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..31231b9ca09 --- /dev/null +++ b/examples/onnx_ptq/petr/prepare_sweep_metadata.py @@ -0,0 +1,179 @@ +# 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 +# +# 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 os + +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", +] +NUM_PREV = 5 +NUM_SWEEPS = 5 + + +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() + + +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"] + ) + + 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_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_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 + ) + 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 + sweep_cam["intrinsics"] = viewpad.astype(np.float32) + sweep_cam["extrinsics"] = lidar2cam_rt.astype(np.float32) + sweep_cam["lidar2img"] = (viewpad @ lidar2cam_rt.T).astype(np.float32) + + for key in ( + "ego2global_translation", + "ego2global_rotation", + "sensor2ego_translation", + "sensor2ego_rotation", + "cam_intrinsic", + ): + sweep_cam.pop(key) + + return sweep_cam + + +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: + 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_lists.append(sweep_cams) + 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 new file mode 100644 index 00000000000..9c4d73286a9 --- /dev/null +++ b/examples/onnx_ptq/petr/quantize.py @@ -0,0 +1,81 @@ +# 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 sys +from pathlib import Path + +from modelopt.onnx.quantization import quantize + +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): + 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=NpzCalibrationReader(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_vovnet_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..86b893c9dee --- /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==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.53.0 +numpy==1.23.5 +nuscenes-devkit==1.1.11 +onnx==1.17.0 +onnx-graphsurgeon==0.6.1 +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==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