384 lines
14 KiB
Python
384 lines
14 KiB
Python
"""
|
|
Utility module for TensorRT engine management and hardware optimization.
|
|
|
|
This module handles the low-level interactions with TensorRT, including:
|
|
- Engine building from ONNX with FP16 conversion and simplification.
|
|
- Persistent timing cache management to speed up subsequent builds.
|
|
- GPU memory allocation and buffer management using Torch-TensorRT interop.
|
|
- Inference execution using CUDA streams and asynchronous execution.
|
|
"""
|
|
|
|
import copy
|
|
import io
|
|
import os
|
|
|
|
import onnx
|
|
import numpy as np
|
|
import tensorrt_rtx as trt
|
|
import torch
|
|
|
|
from collections import OrderedDict
|
|
from logging import error, warning
|
|
from onnxconverter_common import float16
|
|
from onnxsim import simplify
|
|
|
|
# Use polygraphy TensorRTX backend
|
|
os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1"
|
|
|
|
from polygraphy import util
|
|
from polygraphy.backend.common import bytes_from_path
|
|
from polygraphy.backend.trt import ModifyNetworkOutputs, Profile
|
|
from polygraphy.backend.trt import (
|
|
engine_from_bytes,
|
|
engine_from_network,
|
|
network_from_onnx_path,
|
|
network_from_onnx_bytes,
|
|
save_engine,
|
|
)
|
|
from polygraphy.logger import G_LOGGER
|
|
from torch.cuda import nvtx
|
|
from tqdm import tqdm
|
|
|
|
# Path for your persistent hardware profile
|
|
node_dir = os.path.dirname(os.path.realpath(__file__))
|
|
MASTER_CACHE_FILE = f"{node_dir}\\TRT_RTX_ENGINE.cache"
|
|
|
|
G_LOGGER.module_severity = G_LOGGER.ERROR
|
|
|
|
# Map of numpy dtype -> torch dtype
|
|
numpy_to_torch_dtype_dict = {
|
|
np.uint8: torch.uint8,
|
|
np.int8: torch.int8,
|
|
np.int16: torch.int16,
|
|
np.int32: torch.int32,
|
|
np.int64: torch.int64,
|
|
np.float16: torch.float16,
|
|
np.float32: torch.float32,
|
|
np.float64: torch.float64,
|
|
np.complex64: torch.complex64,
|
|
np.complex128: torch.complex128,
|
|
}
|
|
if np.version.full_version >= "1.24.0":
|
|
numpy_to_torch_dtype_dict[np.bool_] = torch.bool
|
|
else:
|
|
numpy_to_torch_dtype_dict[np.bool] = torch.bool
|
|
|
|
# Map of torch dtype -> numpy dtype
|
|
torch_to_numpy_dtype_dict = {
|
|
value: key for (key, value) in numpy_to_torch_dtype_dict.items()
|
|
}
|
|
|
|
def get_timing_cache(config):
|
|
"""
|
|
Retrieves or creates a TensorRT timing cache to accelerate the build process.
|
|
|
|
Args:
|
|
config (trt.IBuilderConfig): The current TensorRT builder configuration.
|
|
|
|
Returns:
|
|
trt.ITimingCache: A timing cache object populated with existing data if found.
|
|
"""
|
|
if os.path.exists(MASTER_CACHE_FILE):
|
|
with open(MASTER_CACHE_FILE, "rb") as f:
|
|
cache_data = f.read()
|
|
return config.create_timing_cache(cache_data)
|
|
return config.create_timing_cache(b"")
|
|
|
|
def save_timing_cache(timing_cache):
|
|
"""
|
|
Serializes and saves the timing cache to disk for future use.
|
|
|
|
Args:
|
|
timing_cache (trt.ITimingCache): The cache containing hardware-specific optimizations.
|
|
"""
|
|
with open(MASTER_CACHE_FILE, "wb") as f:
|
|
f.write(timing_cache.serialize())
|
|
|
|
class TQDMProgressMonitor(trt.IProgressMonitor):
|
|
"""
|
|
Custom TensorRT progress monitor that integrates with TQDM for CLI feedback.
|
|
|
|
Attributes:
|
|
_active_phases (dict): Tracks nested build phases and their progress bars.
|
|
"""
|
|
def __init__(self):
|
|
trt.IProgressMonitor.__init__(self)
|
|
self._active_phases = {}
|
|
self._step_result = True
|
|
self.max_indent = 5
|
|
|
|
def phase_start(self, phase_name, parent_phase, num_steps):
|
|
"""Called by TensorRT when a new build phase starts."""
|
|
leave = False
|
|
try:
|
|
if parent_phase is not None:
|
|
nbIndents = (
|
|
self._active_phases.get(parent_phase, {}).get(
|
|
"nbIndents", self.max_indent
|
|
)
|
|
+ 1
|
|
)
|
|
if nbIndents >= self.max_indent:
|
|
return
|
|
else:
|
|
nbIndents = 0
|
|
leave = True
|
|
self._active_phases[phase_name] = {
|
|
"tq": tqdm(
|
|
total=num_steps, desc=phase_name, leave=leave, position=nbIndents
|
|
),
|
|
"nbIndents": nbIndents,
|
|
"parent_phase": parent_phase,
|
|
}
|
|
except KeyboardInterrupt:
|
|
# The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete.
|
|
_step_result = False
|
|
|
|
def phase_finish(self, phase_name):
|
|
"""Called by TensorRT when a build phase completes."""
|
|
try:
|
|
if phase_name in self._active_phases.keys():
|
|
self._active_phases[phase_name]["tq"].update(
|
|
self._active_phases[phase_name]["tq"].total
|
|
- self._active_phases[phase_name]["tq"].n
|
|
)
|
|
|
|
parent_phase = self._active_phases[phase_name].get("parent_phase", None)
|
|
while parent_phase is not None:
|
|
self._active_phases[parent_phase]["tq"].refresh()
|
|
parent_phase = self._active_phases[parent_phase].get(
|
|
"parent_phase", None
|
|
)
|
|
if (
|
|
self._active_phases[phase_name]["parent_phase"]
|
|
in self._active_phases.keys()
|
|
):
|
|
self._active_phases[
|
|
self._active_phases[phase_name]["parent_phase"]
|
|
]["tq"].refresh()
|
|
del self._active_phases[phase_name]
|
|
pass
|
|
except KeyboardInterrupt:
|
|
_step_result = False
|
|
|
|
def step_complete(self, phase_name, step):
|
|
"""Updates the progress bar for the current step."""
|
|
try:
|
|
if phase_name in self._active_phases.keys():
|
|
self._active_phases[phase_name]["tq"].update(
|
|
step - self._active_phases[phase_name]["tq"].n
|
|
)
|
|
return self._step_result
|
|
except KeyboardInterrupt:
|
|
# There is no need to propagate this exception to TensorRT. We can simply cancel the build.
|
|
return False
|
|
|
|
|
|
class Engine:
|
|
"""
|
|
Wrapper class for TensorRT engine building, loading, and inference.
|
|
|
|
Args:
|
|
engine_path (str): File path where the .trt engine is or will be stored.
|
|
"""
|
|
def __init__(self, engine_path):
|
|
self.engine_path = engine_path
|
|
self.engine = None
|
|
self.context = None
|
|
self.buffers = OrderedDict()
|
|
self.tensors = OrderedDict()
|
|
self.cuda_graph_instance = None # cuda graph
|
|
|
|
def __del__(self):
|
|
del self.engine
|
|
del self.context
|
|
del self.buffers
|
|
del self.tensors
|
|
|
|
def reset(self, engine_path=None):
|
|
"""Clears existing execution context and buffers to free memory."""
|
|
# del self.engine
|
|
del self.context
|
|
del self.buffers
|
|
del self.tensors
|
|
# self.engine_path = engine_path
|
|
|
|
self.buffers = OrderedDict()
|
|
self.tensors = OrderedDict()
|
|
self.inputs = {}
|
|
self.outputs = {}
|
|
|
|
def build(
|
|
self,
|
|
onnx_path,
|
|
input_profile=None,
|
|
enable_refit=False,
|
|
enable_preview=False,
|
|
enable_all_tactics=False,
|
|
timing_cache=None,
|
|
update_output_names=None,
|
|
):
|
|
"""
|
|
Compiles a TensorRT engine from an ONNX model.
|
|
|
|
Args:
|
|
onnx_path (str): Path to the input ONNX file.
|
|
input_profile (list, optional): List of dynamic shape profiles (min, opt, max).
|
|
enable_refit (bool): Whether to enable engine refitting without rebuilds.
|
|
update_output_names (list, optional): Names to override default network outputs.
|
|
|
|
Returns:
|
|
int: 0 if successful, 1 otherwise.
|
|
"""
|
|
# Get the dynamic shapes
|
|
p = [Profile()]
|
|
if input_profile:
|
|
p = [Profile() for i in range(len(input_profile))]
|
|
for _p, i_profile in zip(p, input_profile):
|
|
for name, dims in i_profile.items():
|
|
assert len(dims) == 3
|
|
_p.add(name, min=dims[0], opt=dims[1], max=dims[2])
|
|
|
|
config_kwargs = {}
|
|
if not enable_all_tactics:
|
|
config_kwargs["tactic_sources"] = []
|
|
|
|
# Load ONNX file then convert the model to fp16
|
|
# Probably might cause issues if onnx file given is quantized or whatever
|
|
onnx_bytes = onnx.load(onnx_path)
|
|
onnx_precision = onnx_bytes.graph.input[0].type.tensor_type.elem_type
|
|
onnx_simp, check = simplify(float16.convert_float_to_float16(onnx_bytes) if onnx_precision == onnx.TensorProto.FLOAT else onnx_bytes)
|
|
|
|
if not check:
|
|
error(f"Simplified ONNX model: {onnx_path} could not be validated")
|
|
return 1
|
|
|
|
# Keep simplified model in memory to not overwrite the onnx file in models directory
|
|
# We dont want to be destructive
|
|
onnx_buffer = io.BytesIO()
|
|
onnx_buffer.write(onnx_simp.SerializeToString())
|
|
onnx_buffer.seek(0)
|
|
onnx_buffer.name = "onnx_simp_mem.onnx"
|
|
|
|
network = network_from_onnx_bytes(onnx_buffer.getvalue(), flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM])
|
|
|
|
if update_output_names:
|
|
print(f"Updating network outputs to {update_output_names}")
|
|
network = ModifyNetworkOutputs(network, update_output_names)
|
|
|
|
builder = network[0]
|
|
config = builder.create_builder_config()
|
|
config.progress_monitor = TQDMProgressMonitor()
|
|
|
|
#config.set_flag(trt.BuilderFlag.FP16) if fp16 else None
|
|
config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None
|
|
|
|
# Limit workspace (e.g., to 2GB or 4GB depending on your GPU)
|
|
# 1 << 30 is 1GB
|
|
#config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * (1 << 30))
|
|
config.builder_optimization_level = 5
|
|
|
|
# Use Master Cache
|
|
timing_cache = get_timing_cache(config)
|
|
config.set_timing_cache(timing_cache, False)
|
|
|
|
profiles = copy.deepcopy(p)
|
|
for profile in profiles:
|
|
# Last profile is used for set_calibration_profile.
|
|
calib_profile = profile.fill_defaults(network[1]).to_trt(
|
|
builder, network[1]
|
|
)
|
|
config.add_optimization_profile(calib_profile)
|
|
|
|
result = 0
|
|
|
|
try:
|
|
engine = engine_from_network(network, config)
|
|
save_engine(engine, path=self.engine_path)
|
|
except Exception as e:
|
|
error(f"Failed to build engine: {e}")
|
|
result = 1
|
|
|
|
# Update Master Cache with new findings
|
|
save_timing_cache(timing_cache)
|
|
print(f"Updated Master Cache.")
|
|
|
|
return result
|
|
|
|
def load(self):
|
|
"""Loads a serialized engine from the disk into memory."""
|
|
self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
|
|
|
|
def activate(self, reuse_device_memory=None):
|
|
"""
|
|
Creates an execution context for inference.
|
|
|
|
Args:
|
|
reuse_device_memory (optional): Existing device memory to use for the context.
|
|
"""
|
|
if reuse_device_memory:
|
|
self.context = self.engine.create_execution_context_without_device_memory()
|
|
# self.context.device_memory = reuse_device_memory
|
|
else:
|
|
self.context = self.engine.create_execution_context()
|
|
|
|
def allocate_buffers(self, shape_dict=None, device="cuda"):
|
|
"""
|
|
Allocates GPU memory for all input and output tensors.
|
|
|
|
Args:
|
|
shape_dict (dict, optional): Overrides for tensor shapes.
|
|
device (str): Device to allocate on (default 'cuda').
|
|
"""
|
|
nvtx.range_push("allocate_buffers")
|
|
for idx in range(self.engine.num_io_tensors):
|
|
name = self.engine.get_tensor_name(idx)
|
|
binding = self.engine[idx]
|
|
if shape_dict and binding in shape_dict:
|
|
shape = shape_dict[binding]["shape"]
|
|
else:
|
|
shape = self.context.get_tensor_shape(name)
|
|
|
|
dtype = trt.nptype(self.engine.get_tensor_dtype(name))
|
|
if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
|
|
self.context.set_input_shape(name, shape)
|
|
tensor = torch.empty(
|
|
tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]
|
|
).to(device=device)
|
|
self.tensors[binding] = tensor
|
|
nvtx.range_pop()
|
|
|
|
def infer(self, feed_dict, stream, use_cuda_graph=False):
|
|
"""
|
|
Performs asynchronous inference using CUDA streams.
|
|
|
|
Args:
|
|
feed_dict (dict): Dictionary mapping input names to input tensors.
|
|
stream (int): CUDA stream pointer.
|
|
use_cuda_graph (bool): Whether to use CUDA graph capture (not fully implemented).
|
|
|
|
Returns:
|
|
OrderedDict: The dictionary containing output tensors.
|
|
"""
|
|
nvtx.range_push("set_tensors")
|
|
for name, buf in feed_dict.items():
|
|
self.tensors[name].copy_(buf)
|
|
|
|
for name, tensor in self.tensors.items():
|
|
self.context.set_tensor_address(name, tensor.data_ptr())
|
|
nvtx.range_pop()
|
|
nvtx.range_push("execute")
|
|
noerror = self.context.execute_async_v3(stream)
|
|
if not noerror:
|
|
raise ValueError("ERROR: inference failed.")
|
|
nvtx.range_pop()
|
|
return self.tensors
|
|
|
|
def __str__(self):
|
|
out = ""
|
|
for opt_profile in range(self.engine.num_optimization_profiles):
|
|
for binding_idx in range(self.engine.num_io_tensors):
|
|
name = self.engine.get_tensor_name(binding_idx)
|
|
shape = self.engine.get_tensor_shape(name)
|
|
out += f"\t{name} = {shape}\n"
|
|
return out |