Add kwargs support to torch.onnx.export api for dynamic_axes and other args (#336)

* Add dynamic axes support for pytorch onnx exports

* Change to make more generic to accept kwargs for the torch.onnx.export api

Author: markurtz

src/sparseml/pytorch/utils/exporter.py

@@ -176,6 +176,7 @@ def export_onnx(
         opset: int = DEFAULT_ONNX_OPSET,
         disable_bn_fusing: bool = True,
         convert_qat: bool = False,
+        **export_kwargs,
     ):
         """
         Export an onnx file for the current module and for a sample batch.
@@ -197,7 +198,14 @@ def export_onnx(
             the module being exported, the resulting QAT ONNX model will be converted
             to a fully quantized ONNX model using `quantize_torch_qat_export`. Default
             is False.
+        :param export_kwargs: kwargs to be passed as is to the torch.onnx.export api
+            call. Useful to pass in dyanmic_axes, input_names, output_names, etc.
+            See more on the torch.onnx.export api spec in the PyTorch docs:
+            https://pytorch.org/docs/stable/onnx.html
         """
+        if not export_kwargs:
+            export_kwargs = {}
+
         if isinstance(sample_batch, Dict) and not isinstance(
             sample_batch, collections.OrderedDict
         ):
@@ -217,20 +225,26 @@ def export_onnx(
                 sample_batch, self._module, check_feat_lab_inp=False
             )
 
-        input_names = None
-        if isinstance(sample_batch, Tensor):
-            input_names = ["input"]
-        elif isinstance(sample_batch, Dict):
-            input_names = list(sample_batch.keys())
-            sample_batch = tuple([sample_batch[f] for f in input_names])
-        elif isinstance(sample_batch, Iterable):
-            input_names = [
-                "input_{}".format(index) for index, _ in enumerate(iter(sample_batch))
-            ]
-            if isinstance(sample_batch, List):
-                sample_batch = tuple(sample_batch)  # torch.onnx.export requires tuple
+        if "input_names" not in export_kwargs:
+            if isinstance(sample_batch, Tensor):
+                export_kwargs["input_names"] = ["input"]
+            elif isinstance(sample_batch, Dict):
+                export_kwargs["input_names"] = list(sample_batch.keys())
+                sample_batch = tuple(
+                    [sample_batch[f] for f in export_kwargs["input_names"]]
+                )
+            elif isinstance(sample_batch, Iterable):
+                export_kwargs["input_names"] = [
+                    "input_{}".format(index)
+                    for index, _ in enumerate(iter(sample_batch))
+                ]
+                if isinstance(sample_batch, List):
+                    sample_batch = tuple(
+                        sample_batch
+                    )  # torch.onnx.export requires tuple
 
-        output_names = self.get_output_names(out)
+        if "output_names" not in export_kwargs:
+            export_kwargs["output_names"] = self.get_output_names(out)
 
         # disable active quantization observers because they cannot be exported
         disabled_observers = []
@@ -259,11 +273,10 @@ def export_onnx(
             export_module,
             sample_batch,
             onnx_path,
-            input_names=input_names,
-            output_names=output_names,
             strip_doc_string=True,
             verbose=False,
             opset_version=opset,
+            **export_kwargs,
         )
 
         # re-enable disabled quantization observers