File size: 4,289 Bytes
dd0a8f2 | 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 | from __future__ import annotations
import argparse
import json
from pathlib import Path
import torch
import tensorrt as trt
def trt_dtype_to_torch(dtype):
if dtype == trt.DataType.HALF:
return torch.float16
if dtype == trt.DataType.BF16:
return torch.bfloat16
if dtype == trt.DataType.FLOAT:
return torch.float32
if dtype == trt.DataType.INT8:
return torch.int8
raise RuntimeError(f"Unsupported TensorRT dtype: {dtype}")
def make_inputs(engine, context, device):
inputs = {}
outputs = {}
input_names = []
output_names = []
for idx in range(engine.num_io_tensors):
name = engine.get_tensor_name(idx)
shape = tuple(engine.get_tensor_shape(name))
dtype = trt_dtype_to_torch(engine.get_tensor_dtype(name))
if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
context.set_input_shape(name, shape)
input_names.append(name)
if dtype.is_floating_point:
tensor = torch.randn(shape, dtype=dtype, device=device)
else:
tensor = torch.zeros(shape, dtype=dtype, device=device)
inputs[name] = tensor.contiguous()
else:
output_names.append(name)
outputs[name] = torch.empty(shape, dtype=dtype, device=device)
return input_names, output_names, inputs, outputs
def bind(context, names, tensors):
for name in names:
context.set_tensor_address(name, tensors[name].data_ptr())
def run_once(context, stream):
ok = context.execute_async_v3(stream_handle=stream.cuda_stream)
if not ok:
raise RuntimeError("execute_async_v3 returned false")
def timed_loop(fn, stream, warmup, iters):
for _ in range(warmup):
fn()
stream.synchronize()
times = []
for _ in range(iters):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record(stream)
fn()
end.record(stream)
end.synchronize()
times.append(start.elapsed_time(end))
return times
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--engine_path", required=True)
parser.add_argument("--warmup", type=int, default=5)
parser.add_argument("--iters", type=int, default=30)
parser.add_argument("--cuda_graph", action="store_true")
parser.add_argument("--json_path", type=str, default=None)
args = parser.parse_args()
logger = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(logger, "")
runtime = trt.Runtime(logger)
engine_path = Path(args.engine_path)
engine = runtime.deserialize_cuda_engine(engine_path.read_bytes())
if engine is None:
raise RuntimeError(f"Failed to deserialize engine: {engine_path}")
context = engine.create_execution_context()
device = torch.device("cuda")
input_names, output_names, inputs, outputs = make_inputs(engine, context, device)
bind(context, input_names, inputs)
bind(context, output_names, outputs)
stream = torch.cuda.Stream(device=device)
if args.cuda_graph:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, stream=stream):
run_once(context, torch.cuda.current_stream(device))
def fn():
with torch.cuda.stream(stream):
graph.replay()
else:
fn = lambda: run_once(context, stream)
times = timed_loop(fn, stream, args.warmup, args.iters)
result = {
"engine_path": str(engine_path),
"cuda_graph": args.cuda_graph,
"warmup": args.warmup,
"iters": args.iters,
"mean_ms": sum(times) / len(times),
"min_ms": min(times),
"max_ms": max(times),
"times_ms": times,
"inputs": {name: {"shape": list(inputs[name].shape), "dtype": str(inputs[name].dtype)} for name in input_names},
"outputs": {name: {"shape": list(outputs[name].shape), "dtype": str(outputs[name].dtype)} for name in output_names},
}
print(json.dumps(result, indent=2))
if args.json_path:
out = Path(args.json_path)
out.parent.mkdir(parents=True, exist_ok=True)
out.write_text(json.dumps(result, indent=2))
if __name__ == "__main__":
main()
|