Address PETR example review feedback

Co-Authored-By: Codex <codex@openai.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>

Author: ajrasane

.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|

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

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):

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",

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 \