115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
import tensorrt_rtx as trt
|
|
import argparse
|
|
import os
|
|
import onnx
|
|
from pathlib import Path
|
|
from tqdm import tqdm
|
|
|
|
# Path for your persistent hardware profile
|
|
MASTER_CACHE_FILE = "..\\TRT_RTX_ENGINE.cache"
|
|
|
|
def get_timing_cache(config):
|
|
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):
|
|
with open(MASTER_CACHE_FILE, "wb") as f:
|
|
f.write(timing_cache.serialize())
|
|
|
|
|
|
def get_onnx_precision(onnx_path):
|
|
"""Detects if the ONNX model is FP32 or FP16."""
|
|
model = onnx.load(onnx_path)
|
|
# Check the data type of the first input
|
|
dtype = model.graph.input[0].type.tensor_type.elem_type
|
|
if dtype == onnx.TensorProto.FLOAT16:
|
|
return "fp16"
|
|
return "fp32"
|
|
|
|
def build_engines(input_dir, output_dir):
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
builder = trt.Builder(logger)
|
|
|
|
# Limit workspace (e.g., to 2GB or 4GB depending on your GPU)
|
|
# 1 << 30 is 1GB
|
|
vram_limit = 4 * (1 << 30)
|
|
path = Path(input_dir)
|
|
trt_version = trt.__version__
|
|
|
|
# Grab all .onnx files except 'rife'
|
|
onnx_files = list(f for f in path.glob("*.onnx") if "rife" not in f.name.lower())
|
|
|
|
# Progress Bar configuration
|
|
pbar = tqdm(onnx_files, desc="Converting Models", unit="model")
|
|
|
|
if not onnx_files:
|
|
print(f"No ONNX files found in {input_dir}")
|
|
return
|
|
|
|
for onnx_path in pbar:
|
|
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
|
|
parser = trt.OnnxParser(network, logger)
|
|
config = builder.create_builder_config()
|
|
precision = get_onnx_precision(onnx_path)
|
|
|
|
|
|
# Enforce tactic limits and FP16
|
|
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, vram_limit)
|
|
config.builder_optimization_level = 5
|
|
|
|
if builder.platform_has_fast_fp16:
|
|
precision = "fp16"
|
|
config.set_flag(trt.BuilderFlag.FP16)
|
|
|
|
|
|
# Use Master Cache
|
|
timing_cache = get_timing_cache(config)
|
|
config.set_timing_cache(timing_cache, False)
|
|
|
|
# Parse ONNX
|
|
with open(onnx_path, 'rb') as model:
|
|
if not parser.parse(model.read()):
|
|
print(f"Error parsing {onnx_path}")
|
|
continue
|
|
|
|
# Dynamic Shape Profile
|
|
profile = builder.create_optimization_profile()
|
|
input_name = network.get_input(0).name
|
|
min_shape = (1, 3, 256, 256)
|
|
opt_shape = (1, 3, 512, 512)
|
|
max_shape = (1, 3, 1280, 1280)
|
|
|
|
profile.set_shape(input_name, min_shape, opt_shape, max_shape)
|
|
config.add_optimization_profile(profile)
|
|
|
|
# Filename: ONNXFILE_MIN_OPT_MAX_VERSION.trt
|
|
def fmt(s): return "x".join(map(str, s))
|
|
base_name = onnx_path.stem
|
|
engine_name = f"{base_name}_{fmt(min_shape)}_{fmt(opt_shape)}_{fmt(max_shape)}_{trt_version}.trt"
|
|
engine_path = f"{output_dir}\\{engine_name}"
|
|
|
|
# Build serialized engine
|
|
print(f"Building {engine_path}...")
|
|
serialized_engine = builder.build_serialized_network(network, config)
|
|
|
|
# Update Master Cache with new findings
|
|
save_timing_cache(timing_cache)
|
|
print(f"Updated Master Cache.")
|
|
|
|
if serialized_engine:
|
|
with open(engine_path, "wb") as f:
|
|
f.write(serialized_engine)
|
|
print(f"Saved {engine_path} and updated Master Cache.")
|
|
else:
|
|
print(f"FAILED TO SAVE {engine_path}...")
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("onnx_dir", help="Directory containing .onnx files")
|
|
parser.add_argument("trt_dir", help="Directory to put .trt files")
|
|
args = parser.parse_args()
|
|
build_engines(args.onnx_dir, args.trt_dir)
|