One config for onnx

Author: knowicki-nvidia

CHANGELOG.md

@@ -28,6 +28,7 @@ limitations under the License.
 - change: Install the TensorRT package for architectures other than x86_64
 - change: Disable conversion fallback for TensorRT paths and expose control option in custom config
 - change: Use torch.export.save for Torch-TRT model serialization
+- change: Added export_engine to OnnxConfig for improved export control
 - fix: Correctness command relative tolerance formula
 - fix: Memory management during export and conversion process for Torch
 

examples/09_custom_configurations_for_optimize/optimize.py

@@ -64,12 +64,17 @@ def get_onnx_config():
     - PyTorch -> ONNX
     - TorchScript Script -> ONNX
     - TorchScript Trace -> ONNX
+    - Torch Dynamo -> ONNX
 
     From each ONNX different TensorRT plan will be produced.
     """
     return nav.OnnxConfig(
         opset=17,
         onnx_extended_conversion=True,
+        export_engine=[
+            nav.OnnxTraceExportConfig(),
+            nav.OnnxDynamoExportConfig(),
+        ],
     )
 
 

model_navigator/__init__.py

@@ -23,6 +23,8 @@
     MaxThroughputWithLatencyBudgetStrategy,
     MinLatencyStrategy,
     OnnxConfig,
+    OnnxDynamoExportConfig,
+    OnnxTraceExportConfig,
     OptimizationProfile,
     SelectedRuntimeStrategy,
     TensorFlowConfig,

model_navigator/commands/convert/base.py

@@ -91,7 +91,7 @@ def _execute_single_conversion(
         max_batch_size = cls._get_conversion_max_batch_sizes(device_max_batch_size, dataloader_max_batch_size)
 
         LOGGER.info(
-            """Search for maximal batch size disable. Execute single conversion with max_batch_size %d. """
+            """Search for maximal batch size disable. Execute single conversion with max_batch_size {}. """
             """In order to execute fallback strategy provide `conversion_fallback=True` in custom configuration.""",
             max_batch_size,
         )

model_navigator/commands/export/exporters/__init__.py

@@ -22,7 +22,6 @@
 
 if is_torch2_available():
     from . import (
-        torch2dynamo_onnx,  # noqa: F401
         torch2exportedprogram,  # noqa: F401
     )