HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.37 TensorRT auto-workspace size (#17748)

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Burhan 2024-11-25 05:33:11 -05:00 • committed by GitHub
parent 6a762564c8
commit ee6fde0beb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 24 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

8
ultralytics/engine/exporter.py
View file

@ -781,10 +781,10 @@ class Exporter:
# Engine builder
builder = trt.Builder(logger)
config = builder.create_builder_config()
workspace = int(self.args.workspace * (1 << 30))
if is_trt10:
workspace = int(self.args.workspace * (1 << 30)) if self.args.workspace is not None else 0
if is_trt10 and workspace > 0:
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace)
else: # TensorRT versions 7, 8
elif workspace > 0 and not is_trt10: # TensorRT versions 7, 8
config.max_workspace_size = workspace
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(flag)
@ -823,7 +823,7 @@ class Exporter:
LOGGER.warning(f"{prefix} WARNING ⚠️ 'dynamic=True' model requires max batch size, i.e. 'batch=16'")
profile = builder.create_optimization_profile()
min_shape = (1, shape[1], 32, 32) # minimum input shape
max_shape = (*shape[:2], *(int(max(1, self.args.workspace) * d) for d in shape[2:])) # max input shape
max_shape = (*shape[:2], *(int(max(1, workspace) * d) for d in shape[2:])) # max input shape
for inp in inputs:
profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
config.add_optimization_profile(profile)