Переглянути джерело

Fix TRT `max_workspace_size` deprecation notice (#6856)

* Fix TRT `max_workspace_size` deprecation notice

* Update export.py

* Update export.py
modifyDataloader
Glenn Jocher GitHub 2 роки тому
джерело
коміт
acc58c1dcf
Не вдалося знайти GPG ключ що відповідає даному підпису Ідентифікатор GPG ключа: 4AEE18F83AFDEB23
1 змінених файлів з 1 додано та 0 видалено
  1. +1
    -0
      export.py

+ 1
- 0
export.py Переглянути файл

@@ -218,6 +218,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.max_workspace_size = workspace * 1 << 30
# config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice

flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = builder.create_network(flag)

Завантаження…
Відмінити
Зберегти