File size: 3,075 Bytes
ef79ec1 865bf5f ef79ec1 292b725 ef79ec1 865bf5f ef79ec1 a5425e1 865bf5f ef79ec1 4af75b4 ef79ec1 4af75b4 ef79ec1 | 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 | from __future__ import annotations
import argparse, json, os, time
from pathlib import Path
import tensorrt as trt
LOGGER=trt.Logger(trt.Logger.INFO)
def build(onnx_path:Path, engine_path:Path, workspace_gib:int=8, fp16=True):
parse_path = onnx_path
# TensorRT 11 is strongly typed and removed BuilderFlag.FP16. NVIDIA's
# supported migration path is ModelOpt AutoCast before engine building.
if fp16 and not hasattr(trt.BuilderFlag, "FP16"):
import onnx
from modelopt.onnx.autocast import convert_to_mixed_precision
mixed_path = onnx_path.with_name(onnx_path.stem + "_mixed_fp16.onnx")
print(f"AUTOCAST {onnx_path.name} -> {mixed_path.name}", flush=True)
converted = convert_to_mixed_precision(
onnx_path=str(onnx_path), low_precision_type="fp16",
keep_io_types=True, providers=["cuda:0", "cpu"],
init_conversion_max_bytes=512 * 1024 * 1024,
use_standalone_type_inference=True,
)
onnx.save(converted, str(mixed_path))
parse_path = mixed_path
builder=trt.Builder(LOGGER)
flags = 0
if hasattr(trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH'):
flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network=builder.create_network(flags)
parser=trt.OnnxParser(network, LOGGER)
data=parse_path.read_bytes()
if not parser.parse(data):
errors='\n'.join(str(parser.get_error(i)) for i in range(parser.num_errors))
raise RuntimeError(errors)
config=builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gib*(1<<30))
if hasattr(config, 'builder_optimization_level'): config.builder_optimization_level=5
if fp16 and hasattr(trt.BuilderFlag, "FP16"): config.set_flag(trt.BuilderFlag.FP16)
# Engines are intentionally built for the current L4 and not marked HW-compatible.
t=time.time(); serialized=builder.build_serialized_network(network, config)
if serialized is None: raise RuntimeError('TensorRT build returned None')
engine_path.parent.mkdir(parents=True,exist_ok=True); engine_path.write_bytes(serialized)
return {'seconds':round(time.time()-t,3),'bytes':engine_path.stat().st_size,'fp16':fp16}
def main():
ap=argparse.ArgumentParser(); ap.add_argument('--root',default='/tmp/cf-onnx-out'); ap.add_argument('--component',choices=['dinov2','dsine','all'],default='all'); a=ap.parse_args()
root=Path(a.root); out={}; names=['dinov2','dsine'] if a.component=='all' else [a.component]
for name in names:
src=root/f'onnx/{name}/model.onnx'; dst=root/f'engines/l4-sm89/{name}.plan'
print('BUILD',name,flush=True); out[name]=build(src,dst)
import torch
out['environment']={'gpu':torch.cuda.get_device_name(0),'cc':list(torch.cuda.get_device_capability(0)),'cuda':torch.version.cuda,'tensorrt':trt.__version__}
meta=root/f'trt_build_{a.component}_meta.json'; meta.write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()
|