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

Improved float workspace arg for TRT exports (#9407)

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Jiacong Fang 2024-03-29 23:22:32 +08:00 • committed by GitHub
parent 4a7ccba0af
commit 03d0ffd9f5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 2 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/engine/exporter.py
View file

@ -675,9 +675,7 @@ class Exporter:
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.max_workspace_size = self.args.workspace * 1 << 30
# config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice
config.max_workspace_size = int(self.args.workspace * (1 << 30))
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(flag)
parser = trt.OnnxParser(network, logger)