Update onnx dynamo export method

Author: kacper-kleczewski

model_navigator/commands/export/exporters/torch2dynamo_onnx.py

@@ -21,6 +21,7 @@
 import onnx_graphsurgeon as gs
 import torch  # pytype: disable=import-error
 
+from model_navigator.configuration import TensorRTProfile
 from model_navigator.core.dataloader import load_samples
 from model_navigator.core.tensor import TensorMetadata
 
@@ -37,6 +38,7 @@ def get_model() -> torch.nn.Module:
 def export(
     exported_model_path: str,
     input_metadata: Dict[str, Any],
+    dataloader_trt_profile: Dict[str, Any],
     input_names: List[str],
     output_names: List[str],
     batch_dim: Optional[int],
@@ -45,12 +47,15 @@ def export(
     verbose: bool,
     custom_args: Dict[str, Any],
     navigator_workspace: Optional[str] = None,
+    dataloader_max_batch_size: Optional[int] = None,
+    device_max_batch_size: Optional[int] = None,
 ):
     """Export Torch model using dynamo.
 
     Args:
         exported_model_path (str): Output ONNX model path.
         input_metadata (Dict[str, Any]): List of input metadata.
+        dataloader_trt_profile: Profiles generated based on shapes.
         input_names (List[str]): List of model input names.
         output_names (List[str]): List of model output names.
         batch_dim (Optional[int]): Batch dimension.
@@ -61,6 +66,8 @@ def export(
             When None use current workdir. Defaults to None.
         custom_args (Optional[Dict[str, Any]], optional): Passthrough parameters for torch.jit.trace
             For available arguments check PyTorch documentation: https://pytorch.org/docs/stable/jit.html#torch.jit.trace
+        dataloader_max_batch_size: Maximum batch size in the dataloader. Defaults to None.
+        device_max_batch_size: Maximum batch size that fits on the device. Defaults to None.
     """
     model = get_model()
 
@@ -71,6 +78,14 @@ def export(
     profiling_sample = load_samples("profiling_sample", navigator_workspace, batch_dim)[0]
     input_metadata = TensorMetadata.from_json(input_metadata)
 
+    def expand_batch_dim(tensor, batch_dim, max_batch_size):
+        if batch_dim is not None and tensor.shape[batch_dim] < max_batch_size:
+            expand_shape = list(tensor.shape)
+            expand_shape[batch_dim] = max_batch_size
+            expanded_tensor = tensor.expand(*expand_shape)
+            return expanded_tensor
+        return tensor
+
     dummy_input = {n: torch.from_numpy(val).to(target_device) for n, val in profiling_sample.items()}
     dummy_input = input_metadata.unflatten_sample(dummy_input, wrap_input=False)
 
@@ -80,23 +95,53 @@ def export(
         dummy_input = (*dummy_input, {})
     *args, kwargs = dummy_input
 
+    # Expand batch_dim of tensors to max_batch_size
+    max_batch_size = device_max_batch_size or dataloader_max_batch_size
+    if max_batch_size is not None:
+        args = tuple(
+            expand_batch_dim(arg, batch_dim, max_batch_size) if isinstance(arg, torch.Tensor) else arg for arg in args
+        )
+        kwargs = {
+            k: expand_batch_dim(v, batch_dim, max_batch_size) if isinstance(v, torch.Tensor) else v
+            for k, v in kwargs.items()
+        }
+
     loglevel = logging.WARNING if verbose else logging.ERROR
-    export_options_kwargs = {}
-    export_options_kwargs["diagnostic_options"] = torch.onnx.DiagnosticOptions(verbosity_level=loglevel)
-    if dynamic_shapes:
-        export_options_kwargs["dynamic_shapes"] = True
-    export_options = torch.onnx.ExportOptions(**export_options_kwargs)
 
     root_logger = logging.getLogger()
     original_loglevel = root_logger.getEffectiveLevel()
     root_logger.setLevel(loglevel)
+
+    # Dynamic shapes support
+
+    # Collect trt profile for min and max shape data
+    # FIXME: Use a common structure for the min/max shapes
+    dataloader_trt_profile = TensorRTProfile.from_dict(dataloader_trt_profile)
+    dynamic_shapes = []
+    for name, spec_ in dataloader_trt_profile.items():
+        tensor_metadata = input_metadata.get(name)
+        if not tensor_metadata:
+            continue
+
+        dynamic_shapes_ = {}
+        if max_batch_size is not None and max_batch_size > 1 and len(tensor_metadata.shape) > 0:
+            dynamic_shapes_[0] = torch.export.Dim("batch", min=1, max=max_batch_size)
+
+        for idx in range(1, len(spec_.min)):
+            if spec_.min[idx] != spec_.max[idx]:
+                dynamic_shapes_[idx] = torch.export.Dim(f"{name}__{idx}", min=spec_.min[idx], max=spec_.max[idx])
+
+        dynamic_shapes.append(dynamic_shapes_)
+
     try:
-        exported_model = torch.onnx.dynamo_export(
+        exported_model = torch.onnx.export(
             model,
-            *args,
+            args=tuple(args),
+            kwargs=kwargs,
             **custom_args,
-            **kwargs,
-            export_options=export_options,
+            dynamo=True,
+            dynamic_shapes=dynamic_shapes,
+            fallback=False,
         )
 
         exported_model_path = pathlib.Path(exported_model_path)

model_navigator/commands/export/torch.py

@@ -322,13 +322,16 @@ def _run(
         path: pathlib.Path,
         input_metadata: TensorMetadata,
         output_metadata: TensorMetadata,
+        dataloader_trt_profile: TensorRTProfile,
         target_device: DeviceKind,
         dynamic_axes: Dict[str, Union[Dict[int, str], List[int]]],
         dynamo_dynamic_shapes: Optional[bool],
         verbose: bool,
         custom_args: Dict[str, Any],
         model: Any = None,
         batch_dim: Optional[int] = None,
+        dataloader_max_batch_size: Optional[int] = None,
+        device_max_batch_size: Optional[int] = None,
     ) -> CommandOutput:
         """Execute command.
 
@@ -338,6 +341,7 @@ def _run(
             opset: ONNX opset
             input_metadata: Model inputs metadata
             output_metadata: Model outputs metadata
+            dataloader_trt_profile: Profile from dataloader
             target_device: Target device for export - determine the exported model
             dynamic_axes: Definition of model inputs dynamic axes
             dynamo_dynamic_shapes: Enable dynamo dynamic shapes
@@ -347,6 +351,8 @@ def _run(
             custom_args (Optional[Dict[str, Any]], optional): Passthrough parameters for torch.onnx.dynamo_export
                 Can be used to pass ExportOptions object.
                 For available arguments check PyTorch documentation: https://pytorch.org/docs/stable/onnx.html#torch.onnx.export
+            dataloader_max_batch_size: The maximal batch size obtained from datalaoder
+            device_max_batch_size: The maximal batch size obtained for device
 
         Returns:
             CommandOutput object with status
@@ -394,8 +400,10 @@ def on_exit():
                 "navigator_workspace": workspace.path.as_posix(),
                 "custom_args": custom_args,
                 "verbose": verbose,
+                "dataloader_max_batch_size": dataloader_max_batch_size,
+                "device_max_batch_size": device_max_batch_size,
+                "dataloader_trt_profile": dataloader_trt_profile.to_dict(),
             }
-
             args = parse_kwargs_to_cmd(kwargs)
 
             context.execute_python_script(

model_navigator/pipelines/builders/torch.py

@@ -30,6 +30,7 @@
 from model_navigator.configuration.model.model_config import ModelConfig, ONNXModelConfig
 from model_navigator.pipelines.constants import (
     PIPELINE_TORCH_CONVERSION,
+    PIPELINE_TORCH_DYNAMO_ONNX,
     PIPELINE_TORCH_EXPORT,
     PIPELINE_TORCH_EXPORTEDPROGRAM,
 )
@@ -58,10 +59,7 @@ def torch_export_builder(config: CommonConfig, models_config: Dict[Format, List[
                     ExecutionUnit(command=CopyModelFromPath, model_config=models_config[Format.ONNX][0])
                 )
             else:
-                if model_cfg.dynamo_export:  # pytype: disable=attribute-error
-                    execution_units.append(ExecutionUnit(command=ExportTorch2DynamoONNX, model_config=model_cfg))
-                else:
-                    execution_units.append(ExecutionUnit(command=ExportTorch2ONNX, model_config=model_cfg))
+                execution_units.append(ExecutionUnit(command=ExportTorch2ONNX, model_config=model_cfg))
 
             assert isinstance(model_cfg, ONNXModelConfig)
             if model_cfg.graph_surgeon_optimization:
@@ -70,6 +68,24 @@ def torch_export_builder(config: CommonConfig, models_config: Dict[Format, List[
     return Pipeline(name=PIPELINE_TORCH_EXPORT, execution_units=execution_units)
 
 
+def torch_dynamo_onnx_builder(config: CommonConfig, models_config: Dict[Format, List[ModelConfig]]) -> Pipeline:
+    """Prepare export steps for pipeline.
+
+    Args:
+        config: A configuration for pipelines
+        models_config: List of model configs per format
+
+    Returns:
+        Pipeline with steps for export
+    """
+    execution_units: List[ExecutionUnit] = []
+    for model_cfg in models_config.get(Format.ONNX, []):
+        if model_cfg.dynamo_export:  # pytype: disable=attribute-error
+            execution_units.append(ExecutionUnit(command=ExportTorch2DynamoONNX, model_config=model_cfg))
+
+    return Pipeline(name=PIPELINE_TORCH_DYNAMO_ONNX, execution_units=execution_units)
+
+
 def torch_exportedprogram_builder(config: CommonConfig, models_config: Dict[Format, List[ModelConfig]]) -> Pipeline:
     """Prepare export steps for pipeline.
 

model_navigator/pipelines/constants.py

@@ -27,4 +27,5 @@
 PIPELINE_TORCH_CONVERSION = "PyTorch Conversion"
 PIPELINE_TORCH_EXPORT = "PyTorch Export"
 PIPELINE_TORCH_EXPORTEDPROGRAM = "PyTorch ExportedProgram Export"
+PIPELINE_TORCH_DYNAMO_ONNX = "PyTorch Dynamo ONNX Export"
 PIPELINE_VERIFY_MODELS = "Verify Models"

model_navigator/torch/__init__.py

@@ -45,7 +45,7 @@
     verify_builder,
 )
 from model_navigator.pipelines.builders.find_device_max_batch_size import find_device_max_batch_size_builder
-from model_navigator.pipelines.builders.torch import torch_exportedprogram_builder
+from model_navigator.pipelines.builders.torch import torch_dynamo_onnx_builder, torch_exportedprogram_builder
 from model_navigator.pipelines.wrappers.optimize import optimize_pipeline
 from model_navigator.runners.base import NavigatorRunner
 from model_navigator.runners.utils import default_runners, filter_runners
@@ -140,6 +140,7 @@ def optimize(
         torch_export_builder,
         find_device_max_batch_size_builder,
         torch_exportedprogram_builder,
+        torch_dynamo_onnx_builder,
         torch_conversion_builder,
         torch_tensorrt_conversion_builder,
         tensorrt_conversion_builder,