File size: 3,011 Bytes
688a201 7fb3f7f 6d3559d 688a201 7fb3f7f 688a201 7fb3f7f 2255002 7fb3f7f 688a201 7fb3f7f 2255002 7fb3f7f 688a201 7fb3f7f 688a201 7fb3f7f 688a201 | 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 | import json
import os
import sys
from pathlib import Path
import tensorrt as trt
from huggingface_hub import snapshot_download
ROOT = Path(snapshot_download(
os.environ.get('CF_RUNTIME_REPO','patdev/Companion-Forge-L4-ONNX'),
repo_type='model', token=os.environ.get('HF_TOKEN'), local_dir='/tmp/cf-custom',
allow_patterns=['onnx/custom-ops/*.onnx','plugins/tensorrt/companion_sparse_trt.py'],
))
sys.path.insert(0, str(ROOT/'plugins/tensorrt'))
import companion_sparse_trt # noqa: F401,E402
LOGGER = trt.Logger(trt.Logger.INFO)
PROFILES = {
'SparseSubdivide': {'feats': (32,8), 'coords': (32,4)},
'SparseDownsample': {'feats': (64,8), 'coords': (64,4)},
'SparseUpsample': {'feats': (8,8), 'target_coords': (64,4), 'inverse': (64,)},
'SparseWindowAttention': {'qkv': (64,3,4,16), 'coords': (64,4)},
'SparseConv3D': {'feats': (64,8), 'coords': (64,4), 'weight': (16,3,3,3,8), 'bias': (16,)},
'MeshTopologyExtract': {
'verts_grid': (274625,3), 'sdf': (274625,), 'cube_idx': (262144,8),
'beta': (262144,12), 'alpha': (262144,8), 'gamma': (262144,), 'colors_grid': (274625,6),
},
}
def build(name: str):
onnx_path = ROOT/'onnx/custom-ops'/f'{name}.onnx'
builder = trt.Builder(LOGGER)
network = builder.create_network(0)
parser = trt.OnnxParser(network, LOGGER)
if not parser.parse(onnx_path.read_bytes()):
errs = [str(parser.get_error(i)) for i in range(parser.num_errors)]
raise RuntimeError(f'{name} parse failed:\n' + '\n'.join(errs))
cfg = builder.create_builder_config()
cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)
prof = builder.create_optimization_profile()
input_report = {}
dynamic = False
for i in range(network.num_inputs):
inp = network.get_input(i)
shape = PROFILES[name].get(inp.name)
if shape is None:
raise KeyError(f'{name}: missing profile for {inp.name}')
input_report[inp.name] = tuple(inp.shape)
if any(int(d) < 0 for d in inp.shape):
dynamic = True
prof.set_shape(inp.name, shape, shape, shape)
if dynamic:
cfg.add_optimization_profile(prof)
blob = builder.build_serialized_network(network, cfg)
if blob is None:
raise RuntimeError(f'{name} build_serialized_network returned None')
out = Path('/tmp/custom-engines')/f'{name}.plan'
out.parent.mkdir(parents=True, exist_ok=True)
out.write_bytes(bytes(blob))
return {
'op': name,
'onnx_bytes': onnx_path.stat().st_size,
'plan_bytes': out.stat().st_size,
'inputs': input_report,
'outputs': [network.get_output(i).name for i in range(network.num_outputs)],
}
if __name__ == '__main__':
report=[]
for name in PROFILES:
print(f'=== BUILD {name} ===', flush=True)
r=build(name); print('OK',json.dumps(r),flush=True); report.append(r)
Path('/tmp/custom-engines/report.json').write_text(json.dumps(report,indent=2))
|