(
engine_path: Path,
image: torch.Tensor,
depth: torch.Tensor,
torch_output: np.ndarray,
args: argparse.Namespace,
)
| 423 | |
| 424 | |
| 425 | def run_tensorrt_engine( |
| 426 | engine_path: Path, |
| 427 | image: torch.Tensor, |
| 428 | depth: torch.Tensor, |
| 429 | torch_output: np.ndarray, |
| 430 | args: argparse.Namespace, |
| 431 | ) -> dict[str, Any]: |
| 432 | logger = trt.Logger(trt.Logger.WARNING) |
| 433 | runtime = trt.Runtime(logger) |
| 434 | engine = runtime.deserialize_cuda_engine(engine_path.read_bytes()) |
| 435 | if engine is None: |
| 436 | return {"ok": False, "stage": "deserialize"} |
| 437 | |
| 438 | context = engine.create_execution_context() |
| 439 | output_shape = tuple(engine.get_tensor_shape("depth_refined")) |
| 440 | output_dtype = trt_dtype_to_torch(engine.get_tensor_dtype("depth_refined")) |
| 441 | output = torch.empty(output_shape, dtype=output_dtype, device=image.device).contiguous() |
| 442 | |
| 443 | context.set_tensor_address("image", int(image.data_ptr())) |
| 444 | context.set_tensor_address("depth", int(depth.data_ptr())) |
| 445 | context.set_tensor_address("depth_refined", int(output.data_ptr())) |
| 446 | |
| 447 | if image.device.type == "cuda": |
| 448 | torch.cuda.synchronize() |
| 449 | stream = torch.cuda.Stream() |
| 450 | with torch.cuda.stream(stream): |
| 451 | for _ in range(args.trt_warmup): |
| 452 | if not context.execute_async_v3(stream.cuda_stream): |
| 453 | return {"ok": False, "stage": "warmup_execute"} |
| 454 | stream.synchronize() |
| 455 | |
| 456 | start = time.perf_counter() |
| 457 | with torch.cuda.stream(stream): |
| 458 | for _ in range(args.trt_runs): |
| 459 | if not context.execute_async_v3(stream.cuda_stream): |
| 460 | return {"ok": False, "stage": "execute"} |
| 461 | stream.synchronize() |
| 462 | infer_s = (time.perf_counter() - start) / args.trt_runs |
| 463 | else: |
| 464 | start = time.perf_counter() |
| 465 | for _ in range(args.trt_runs): |
| 466 | if not context.execute_v2([]): |
| 467 | return {"ok": False, "stage": "execute_v2"} |
| 468 | infer_s = (time.perf_counter() - start) / args.trt_runs |
| 469 | |
| 470 | trt_output = output.detach().cpu().numpy() |
| 471 | return { |
| 472 | "ok": True, |
| 473 | "engine": str(engine_path), |
| 474 | "runs": args.trt_runs, |
| 475 | "warmup": args.trt_warmup, |
| 476 | "infer_s_avg": infer_s, |
| 477 | "depth": summarize_depth(trt_output.squeeze(0)), |
| 478 | "vs_torch": compare_outputs(torch_output, trt_output), |
| 479 | } |
| 480 | |
| 481 | |
| 482 | def trt_dtype_to_torch(dtype: trt.DataType) -> torch.dtype: |
no test coverage detected