Spaces:
Configuration error
Configuration error
File size: 4,719 Bytes
3624d0b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Build/cache the RCLane TensorRT FP16 engine and verify its outputs.
TensorRT is used through ONNX Runtime's TensorRT Execution Provider so the
cached ``.engine`` remains consumable by the same runtime pipeline. The script
refuses silent provider fallback and records numerical/performance comparisons
against CUDA FP32.
"""
import argparse
import json
import os
import time
from pathlib import Path
import numpy as np
from test_video_onnx import OUTPUT_NAMES, create_session, timing_summary
def parse_args():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True)
parser.add_argument("--cache-dir", default="exports/trt_cache")
parser.add_argument("--report", default=None)
parser.add_argument("--warmup", type=int, default=10)
parser.add_argument("--iterations", type=int, default=100)
parser.add_argument("--seed", type=int, default=0)
return parser.parse_args()
def benchmark(session, images, warmup, iterations):
for _ in range(warmup):
session.run(list(OUTPUT_NAMES), {"images": images})
timings = []
outputs = None
for _ in range(iterations):
started = time.perf_counter()
outputs = session.run(list(OUTPUT_NAMES), {"images": images})
timings.append((time.perf_counter() - started) * 1000.0)
return outputs, timing_summary(timings)
def main():
args = parse_args()
if args.warmup < 0 or args.iterations <= 0:
raise ValueError("warmup must be non-negative and iterations positive")
model_path = Path(args.model).expanduser().resolve()
cache_dir = Path(args.cache_dir).expanduser().resolve()
report_path = (
Path(args.report).expanduser().resolve()
if args.report else cache_dir / "build_report.json"
)
if not model_path.is_file():
raise FileNotFoundError(model_path)
cache_dir.mkdir(parents=True, exist_ok=True)
report_path.parent.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(args.seed)
images = rng.normal(size=(1, 3, 320, 800)).astype(np.float32)
build_started = time.perf_counter()
trt_session = create_session(
model_path, "tensorrt", allow_tf32=False, trt_cache_dir=cache_dir
)
build_seconds = time.perf_counter() - build_started
if trt_session.get_providers()[0] != "TensorrtExecutionProvider":
raise RuntimeError("TensorRT provider silently fell back")
cuda_session = create_session(model_path, "cuda", allow_tf32=False)
trt_outputs, trt_timing = benchmark(
trt_session, images, args.warmup, args.iterations
)
cuda_outputs, cuda_timing = benchmark(
cuda_session, images, args.warmup, args.iterations
)
comparisons = {}
for name, trt_output, cuda_output in zip(
OUTPUT_NAMES, trt_outputs, cuda_outputs
):
difference = trt_output.astype(np.float64) - cuda_output.astype(
np.float64
)
comparisons[name] = {
"shape": list(trt_output.shape),
"max_abs": float(np.max(np.abs(difference))),
"mean_abs": float(np.mean(np.abs(difference))),
"rmse": float(np.sqrt(np.mean(difference ** 2))),
}
cache_files = []
for path in sorted(cache_dir.iterdir()):
if path.is_file():
cache_files.append({
"path": str(path),
"size_bytes": path.stat().st_size,
})
engines = [
item for item in cache_files if item["path"].endswith(".engine")
]
if not engines:
raise RuntimeError(f"TensorRT did not create an engine in {cache_dir}")
report = {
"model": str(model_path),
"provider": trt_session.get_providers()[0],
"precision": "fp16",
"input_shape": list(images.shape),
"build_or_cache_load_seconds": build_seconds,
"tensorrt_timing": trt_timing,
"cuda_fp32_timing": cuda_timing,
"speedup_from_median": (
cuda_timing["median_ms"] / trt_timing["median_ms"]
),
"output_comparison_to_cuda_fp32": comparisons,
"cache_files": cache_files,
}
temporary = report_path.with_suffix(report_path.suffix + ".tmp")
with temporary.open("w") as handle:
json.dump(report, handle, indent=2)
handle.write("\n")
os.replace(temporary, report_path)
print(f"TensorRT engine OK: {engines[0]['path']}")
print(f"report: {report_path}")
print(
"median inference: TensorRT={:.3f}ms CUDA={:.3f}ms speedup={:.2f}x".format(
trt_timing["median_ms"], cuda_timing["median_ms"],
report["speedup_from_median"],
)
)
if __name__ == "__main__":
main()
|