| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import tensorrt as trt |
|
|
|
|
| def _parse_enum_list(value: str | None, enum_cls): |
| if not value: |
| return [] |
| out = [] |
| for raw in value.split(","): |
| name = raw.strip().upper() |
| if not name: |
| continue |
| if not hasattr(enum_cls, name): |
| valid = ", ".join(x for x in dir(enum_cls) if x.isupper()) |
| raise ValueError(f"Unknown {enum_cls.__name__} value {raw!r}; valid: {valid}") |
| out.append(getattr(enum_cls, name)) |
| return out |
|
|
|
|
| def _enum_choice(value: str | None, enum_cls): |
| if value is None: |
| return None |
| name = value.strip().upper() |
| if not hasattr(enum_cls, name): |
| valid = ", ".join(x.lower() for x in dir(enum_cls) if x.isupper()) |
| raise ValueError(f"Unknown {enum_cls.__name__} value {value!r}; valid: {valid}") |
| return getattr(enum_cls, name) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--onnx_path", required=True) |
| parser.add_argument("--engine_path", required=True) |
| parser.add_argument("--precision", choices=("none", "fp16", "bf16", "bf16_fp16"), default="bf16_fp16") |
| parser.add_argument("--workspace_gb", type=float, default=64.0) |
| parser.add_argument("--builder_optimization_level", type=int, default=5) |
| parser.add_argument("--max_aux_streams", type=int, default=None) |
| parser.add_argument("--strongly_typed", action="store_true") |
| parser.add_argument("--sparse_weights", action="store_true") |
| parser.add_argument("--prefer_precision_constraints", action="store_true") |
| parser.add_argument("--obey_precision_constraints", action="store_true") |
| parser.add_argument("--direct_io", action="store_true") |
| parser.add_argument("--no_tf32", action="store_true") |
| parser.add_argument("--tactic_sources", type=str, default=None, |
| help="Comma-separated TensorRT tactic sources, e.g. CUBLAS_LT,CUDNN,JIT_CONVOLUTIONS") |
| parser.add_argument("--preview_features", type=str, default=None, |
| help="Comma-separated TensorRT preview features to enable.") |
| parser.add_argument("--tiling_optimization_level", type=str, default=None, |
| help="TensorRT tiling optimization level: none, fast, moderate, or full.") |
| parser.add_argument("--hardware_compatibility_level", type=str, default=None, |
| help="TensorRT hardware compatibility level: none, ampere_plus, same_compute_capability.") |
| parser.add_argument("--profiling_verbosity", type=str, default=None, |
| help="TensorRT profiling verbosity: none, layer_names_only, detailed.") |
| parser.add_argument("--timing_cache", type=str, default=None, |
| help="Target-specific TensorRT timing-cache file to read and update.") |
| parser.add_argument("--version_compatible", action="store_true", |
| help="Build a version-compatible plan. Native exact-version plans remain the default.") |
| parser.add_argument("--exclude_lean_runtime", action="store_true", |
| help="Exclude the lean runtime from a version-compatible plan.") |
| args = parser.parse_args() |
|
|
| onnx_path = Path(args.onnx_path) |
| engine_path = Path(args.engine_path) |
| engine_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| logger = trt.Logger(trt.Logger.INFO) |
| trt.init_libnvinfer_plugins(logger, "") |
| builder = trt.Builder(logger) |
| network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) |
| if args.strongly_typed: |
| network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) |
| network = builder.create_network(network_flags) |
| parser = trt.OnnxParser(network, logger) |
| if not parser.parse_from_file(str(onnx_path)): |
| for idx in range(parser.num_errors): |
| print(parser.get_error(idx)) |
| raise RuntimeError(f"Failed to parse ONNX: {onnx_path}") |
|
|
| config = builder.create_builder_config() |
| config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(args.workspace_gb * (1024 ** 3))) |
| if hasattr(config, "builder_optimization_level"): |
| config.builder_optimization_level = args.builder_optimization_level |
| if args.max_aux_streams is not None and hasattr(config, "max_aux_streams"): |
| config.max_aux_streams = args.max_aux_streams |
| if not args.strongly_typed: |
| if args.precision in {"fp16", "bf16_fp16"}: |
| config.set_flag(trt.BuilderFlag.FP16) |
| if args.precision in {"bf16", "bf16_fp16"}: |
| config.set_flag(trt.BuilderFlag.BF16) |
| if args.sparse_weights: |
| config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS) |
| if args.prefer_precision_constraints: |
| config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS) |
| if args.obey_precision_constraints: |
| config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) |
| if args.direct_io: |
| config.set_flag(trt.BuilderFlag.DIRECT_IO) |
| if args.no_tf32 and hasattr(config, "clear_flag"): |
| config.clear_flag(trt.BuilderFlag.TF32) |
| tactic_sources = _parse_enum_list(args.tactic_sources, trt.TacticSource) |
| if tactic_sources: |
| mask = 0 |
| for source in tactic_sources: |
| mask |= 1 << int(source) |
| config.set_tactic_sources(mask) |
| preview_features = _parse_enum_list(args.preview_features, trt.PreviewFeature) |
| for feature in preview_features: |
| config.set_preview_feature(feature, True) |
| tiling_level = _enum_choice(args.tiling_optimization_level, trt.TilingOptimizationLevel) |
| if tiling_level is not None and hasattr(config, "tiling_optimization_level"): |
| config.tiling_optimization_level = tiling_level |
| hardware_level = _enum_choice(args.hardware_compatibility_level, trt.HardwareCompatibilityLevel) |
| if hardware_level is not None and hasattr(config, "hardware_compatibility_level"): |
| config.hardware_compatibility_level = hardware_level |
| profiling_verbosity = _enum_choice(args.profiling_verbosity, trt.ProfilingVerbosity) |
| if profiling_verbosity is not None: |
| config.profiling_verbosity = profiling_verbosity |
| if args.version_compatible: |
| config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE) |
| if args.exclude_lean_runtime: |
| if not args.version_compatible: |
| raise ValueError("--exclude_lean_runtime requires --version_compatible") |
| config.set_flag(trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME) |
| timing_cache_path = Path(args.timing_cache) if args.timing_cache else None |
| if timing_cache_path is not None: |
| timing_blob = timing_cache_path.read_bytes() if timing_cache_path.is_file() else b"" |
| timing_cache = config.create_timing_cache(timing_blob) |
| if not config.set_timing_cache(timing_cache, ignore_mismatch=False): |
| raise RuntimeError(f"TensorRT rejected timing cache: {timing_cache_path}") |
|
|
| print(f"onnx_path={onnx_path}") |
| print(f"engine_path={engine_path}") |
| print(f"precision={args.precision}") |
| print(f"workspace_gb={args.workspace_gb}") |
| print(f"builder_optimization_level={args.builder_optimization_level}") |
| print(f"max_aux_streams={args.max_aux_streams}") |
| print(f"strongly_typed={args.strongly_typed}") |
| print(f"sparse_weights={args.sparse_weights}") |
| print(f"prefer_precision_constraints={args.prefer_precision_constraints}") |
| print(f"obey_precision_constraints={args.obey_precision_constraints}") |
| print(f"direct_io={args.direct_io}") |
| print(f"no_tf32={args.no_tf32}") |
| print(f"tactic_sources={args.tactic_sources}") |
| print(f"preview_features={args.preview_features}") |
| print(f"tiling_optimization_level={args.tiling_optimization_level}") |
| print(f"hardware_compatibility_level={args.hardware_compatibility_level}") |
| print(f"profiling_verbosity={args.profiling_verbosity}") |
| print(f"timing_cache={args.timing_cache}") |
| print(f"version_compatible={args.version_compatible}") |
| print(f"exclude_lean_runtime={args.exclude_lean_runtime}") |
| print(f"num_inputs={network.num_inputs}") |
| for idx in range(network.num_inputs): |
| tensor = network.get_input(idx) |
| print(f"input[{idx}] name={tensor.name} shape={tuple(tensor.shape)} dtype={tensor.dtype}") |
| print(f"num_outputs={network.num_outputs}") |
| for idx in range(network.num_outputs): |
| tensor = network.get_output(idx) |
| print(f"output[{idx}] name={tensor.name} shape={tuple(tensor.shape)} dtype={tensor.dtype}") |
|
|
| serialized = builder.build_serialized_network(network, config) |
| if serialized is None: |
| raise RuntimeError("TensorRT build_serialized_network returned None") |
| engine_path.write_bytes(bytes(serialized)) |
| print(f"engine_bytes={engine_path.stat().st_size}") |
| if timing_cache_path is not None: |
| timing_cache_path.parent.mkdir(parents=True, exist_ok=True) |
| timing_cache_path.write_bytes(bytes(config.get_timing_cache().serialize())) |
| print(f"timing_cache_bytes={timing_cache_path.stat().st_size}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
