Add argument in the CLI to specify device to do the ONNX export on (#634

)

* add args device

* do validation on device as well

* handle device for list and tuple

* fix test

* Update optimum/exporters/onnx/convert.py

Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>

* Update optimum/exporters/onnx/convert.py

Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>

* fix test

* fix torch import

* fix import2

* debug

* run all

Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>

optimum/commands/export/onnx.py

@@ -45,6 +45,12 @@ def parse_args_onnx(parser):
  "conditional generation. If enabled the encoder and decoder of the model are exported separately."
  ),
  )
+ optional_group.add_argument(
+ "--device",
+ type=str,
+ default="cpu",
+ help='The device to use to do the export. Defaults to "cpu".',
+ )
  optional_group.add_argument(
  "--opset",
  type=int,

optimum/exporters/onnx/__main__.py

@@ -139,10 +139,16 @@ def main():
  output_dir=args.output.parent,
  output_names=output_names,
  input_shapes=input_shapes,
+ device=args.device,
  )
  else:
  onnx_inputs, onnx_outputs = export(
- model=model, config=onnx_config, output=args.output, opset=args.opset, input_shapes=input_shapes
+ model=model,
+ config=onnx_config,
+ output=args.output,
+ opset=args.opset,
+ input_shapes=input_shapes,
+ device=args.device,
  )
 
  try:
@@ -155,6 +161,7 @@ def main():
  atol=args.atol,
  output_dir=args.output.parent,
  output_names=output_names,
+ device=args.device,
  )
  else:
  validate_model_outputs(
@@ -163,6 +170,7 @@ def main():
  onnx_model=args.output,
  onnx_named_outputs=onnx_outputs,
  atol=args.atol,
+ device=args.device,
  )
 
  logger.info(f"The ONNX export succeeded and the exported model was saved at: {args.output.parent.as_posix()}")
@@ -178,7 +186,7 @@ def main():
  )
  except Exception as e:
  logger.error(
- f"An error occured with error {e}.\n The model was exported model was saved at: {args.output.parent.as_posix()}"
+ f"An error occured with the error message: {e}.\n The exported model was saved at: {args.output.parent.as_posix()}"
  )
 
 

optimum/exporters/onnx/convert.py

@@ -28,6 +28,7 @@
 from ...onnx.utils import _get_onnx_external_data_tensors, check_model_uses_external_data
 from ...utils import TORCH_MINIMUM_VERSION, is_diffusers_available, is_torch_onnx_support_available, logging
 from .base import OnnxConfig
+from .utils import recursive_to_device
 
 
 if is_torch_available():
@@ -90,6 +91,7 @@ def validate_models_outputs(
  atol: Optional[float] = None,
  output_names: Optional[List[str]] = None,
  input_shapes: Optional[Dict] = None,
+ device: str = "cpu",
 ):
  """
  Validates the export of several models, by checking that the outputs from both the reference and the exported model match.
@@ -109,6 +111,9 @@ def validate_models_outputs(
  If None, will use the keys from the `models_and_onnx_configs` as names.
  input_shapes (`Optional[Dict]`, defaults to `None`):
  If specified, allows to use specific shapes to validate the ONNX model on.
+ device (`str`, defaults to `"cpu"`):
+ The device on which the ONNX models will be validated. Either `cpu` or `cuda`. Validation on a CUDA device is supported only for PyTorch.
+
  Raises:
  ValueError: If the outputs shapes or values do not match between the reference and the exported model.
  """
@@ -136,6 +141,7 @@ def validate_models_outputs(
  onnx_named_outputs=onnx_named_outputs[i],
  atol=atol,
  input_shapes=input_shapes,
+ device=device,
  )
 
 
@@ -146,6 +152,7 @@ def validate_model_outputs(
  onnx_named_outputs: List[str],
  atol: Optional[float] = None,
  input_shapes: Optional[Dict] = None,
+ device: str = "cpu",
 ):
  """
  Validates the export by checking that the outputs from both the reference and the exported model match.
@@ -163,6 +170,8 @@ def validate_model_outputs(
  The absolute tolerance in terms of outputs difference between the reference and the exported model.
  input_shapes (`Optional[Dict]`, defaults to `None`):
  If specified, allows to use specific shapes to validate the ONNX model on.
+ device (`str`, defaults to `"cpu"`):
+ The device on which the ONNX model will be validated. Either `cpu` or `cuda`. Validation on a CUDA device is supported only for PyTorch.
 
  Raises:
  ValueError: If the outputs shapes or values do not match between the reference and the exported model.
@@ -185,11 +194,20 @@ def validate_model_outputs(
 
  # Create ONNX Runtime session
  options = SessionOptions()
- session = InferenceSession(onnx_model.as_posix(), options, providers=["CPUExecutionProvider"])
+
+ if device.startswith("cuda"):
+ provider = "CUDAExecutionProvider"
+ else:
+ provider = "CPUExecutionProvider"
+
+ session = InferenceSession(onnx_model.as_posix(), options, providers=[provider])
 
  # Compute outputs from the reference model
  if is_torch_available() and isinstance(reference_model, nn.Module):
- reference_model.to("cpu")
+ reference_model.to(device)
+
+ for key, value in reference_model_inputs.items():
+ reference_model_inputs[key] = recursive_to_device(value=value, device=device)
 
  ref_outputs = reference_model(**reference_model_inputs)
  ref_outputs_dict = {}
@@ -214,9 +232,9 @@ def validate_model_outputs(
  for name, value in reference_model_inputs_for_validation.items():
  if isinstance(value, (list, tuple)):
  value = config.flatten_output_collection_property(name, value)
- onnx_inputs.update({tensor_name: pt_tensor.numpy() for tensor_name, pt_tensor in value.items()})
+ onnx_inputs.update({tensor_name: pt_tensor.cpu().numpy() for tensor_name, pt_tensor in value.items()})
  else:
- onnx_inputs[name] = value.numpy()
+ onnx_inputs[name] = value.cpu().numpy()
 
  # Compute outputs from the ONNX model
  onnx_outputs = session.run(onnx_named_outputs, onnx_inputs)
@@ -245,9 +263,9 @@ def validate_model_outputs(
  value_failures = []
  for name, ort_value in zip(onnx_named_outputs, onnx_outputs):
  if is_torch_available() and isinstance(reference_model, nn.Module):
- ref_value = ref_outputs_dict[name].detach().numpy()
+ ref_value = ref_outputs_dict[name].detach().cpu().numpy()
  else:
- ref_value = ref_outputs_dict[name].numpy()
+ ref_value = ref_outputs_dict[name].cpu().numpy()
  logger.info(f'\t- Validating ONNX Model output "{name}":')
 
  # Shape
@@ -296,7 +314,7 @@ def export_pytorch(
  The version of the ONNX operator set to use.
  output (`Path`):
  Directory to store the exported ONNX model.
- device (`str`, *optional*, defaults to `cpu`):
+ device (`str`, defaults to `"cpu"`):
  The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
  export on CUDA devices.
  input_shapes (`optional[Dict]`, defaults to `None`):
@@ -491,7 +509,7 @@ def export_models(
  output_names (`Optional[List[str]]`, defaults to `None`):
  The names to use for the exported ONNX files. The order must be the same as the order of submodels in the ordered dict `models_and_onnx_configs`.
  If None, will use the keys from `models_and_onnx_configs` as names.
- device (`str`, *optional*, defaults to `cpu`):
+ device (`str`, defaults to `"cpu"`):
  The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
  export on CUDA devices.
  input_shapes (`Optional[Dict]`, defaults to `None`):

optimum/exporters/onnx/utils.py

@@ -15,9 +15,10 @@
 """Utility functions."""
 
 import copy
-from typing import TYPE_CHECKING, Dict, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Tuple, Union
 
 import packaging
+import torch
 from transformers.utils import is_tf_available, is_torch_available
 
 from ...utils import ORT_QUANTIZE_MINIMUM_VERSION, TORCH_MINIMUM_VERSION, is_diffusers_available
@@ -174,3 +175,18 @@ def get_stable_diffusion_models_for_export(
  models_for_export["vae"] = (vae, vae_onnx_config)
 
  return models_for_export
+
+
+def recursive_to_device(value: Union[Tuple, List, "torch.Tensor"], device: str):
+ if isinstance(value, tuple):
+ value = list(value)
+ for i, val in enumerate(value):
+ value[i] = recursive_to_device(val, device)
+ value = tuple(value)
+ elif isinstance(value, list):
+ for i, val in enumerate(value):
+ value[i] = recursive_to_device(val, device)
+ elif isinstance(value, torch.Tensor):
+ value = value.to(device)
+
+ return value

tests/exporters/test_exporters_onnx_cli.py

@@ -80,21 +80,18 @@ def _onnx_export(self, test_name: str, model_name: str, task: Optional[str], for
 
  with TemporaryDirectory() as tmpdir:
  for_ort = " --for-ort " if for_ort is True else " "
- try:
- if task is not None:
- subprocess.run(
- f"python3 -m optimum.exporters.onnx --model {model_name}{for_ort}--task {task} {tmpdir}",
- shell=True,
- check=True,
- )
- else:
- subprocess.run(
- f"python3 -m optimum.exporters.onnx --model {model_name}{for_ort}{tmpdir}",
- shell=True,
- check=True,
- )
- except Exception as e:
- self.fail(f"{test_name} raised: {e}")
+ if task is not None:
+ subprocess.run(
+ f"python3 -m optimum.exporters.onnx --model {model_name}{for_ort}--task {task} {tmpdir}",
+ shell=True,
+ check=True,
+ )
+ else:
+ subprocess.run(
+ f"python3 -m optimum.exporters.onnx --model {model_name}{for_ort}{tmpdir}",
+ shell=True,
+ check=True,
+ )
 
  def test_all_models_tested(self):
  # make sure we test all models