#
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
from polygraphy import constants, mod, util
from polygraphy.common import TensorMetadata
from polygraphy.comparator import IterationResult
from polygraphy.comparator.data_loader import DataLoaderCache
from polygraphy.logger import G_LOGGER, LogMode
from polygraphy.tools.args import util as args_util
from polygraphy.tools.args.base import BaseArgs
from polygraphy.tools.args.comparator.data_loader import DataLoaderArgs
from polygraphy.tools.args.model import ModelArgs
from polygraphy.tools.script import Script, make_invocable, make_invocable_if_nondefault_kwargs
onnx_backend = mod.lazy_import("polygraphy.backend.onnx")
onnxrt_backend = mod.lazy_import("polygraphy.backend.onnxrt")
[docs]@mod.export()
class OnnxInferShapesArgs(BaseArgs):
"""
ONNX Shape Inference: ONNX shape inference.
Depends on:
- OnnxLoadArgs
- DataLoaderArgs: if allow_force_fallback == True
"""
def __init__(self, default: bool = None, allow_force_fallback: bool = None):
"""
Args:
default (bool):
Whether shape inference should be enabled by default.
Defaults to False.
allow_force_fallback (bool):
Whether fallback shape inference using ONNX-Runtime should be allowed.
Defaults to False.
"""
super().__init__()
self._default = util.default(default, False)
self._allow_force_fallback = util.default(allow_force_fallback, False)
def add_parser_args_impl(self):
shape_infer_group = self.group.add_mutually_exclusive_group()
if self._default:
shape_infer_group.add_argument(
"--no-shape-inference",
help="Disable ONNX shape inference when loading the model",
dest="do_shape_inference",
action="store_false",
default=True,
)
else:
shape_infer_group.add_argument(
"--shape-inference",
"--do-shape-inference",
help="Enable ONNX shape inference when loading the model",
dest="do_shape_inference",
action="store_true",
default=False,
)
if self._allow_force_fallback:
shape_infer_group.add_argument(
"--force-fallback-shape-inference",
help="Force Polygraphy to use ONNX-Runtime to determine metadata for "
"tensors in the graph. This can be useful in cases where ONNX shape inference does not generate correct information. "
"Note that this will cause dynamic dimensions to become static. ",
action="store_true",
default=None,
)
self.group.add_argument(
"--no-onnxruntime-shape-inference",
help="Disable using ONNX-Runtime's shape inference utilities. "
"This will force Polygraphy to use `onnx.shape_inference` instead. "
"Note that ONNX-Runtime's shape inference utilities may be more performant and memory-efficient. ",
dest="allow_onnxruntime",
action="store_false",
default=None,
)
[docs] def parse_impl(self, args):
"""
Parses command-line arguments and populates the following attributes:
Attributes:
do_shape_inference (bool): Whether to do shape inference.
force_fallback (bool): Whether to force fallback shape inference.
allow_onnxruntime (bool): Whether to allow ONNX-Runtime's shape inference utilities to be used.
"""
self.do_shape_inference = args_util.get(args, "do_shape_inference")
self.force_fallback = args_util.get(args, "force_fallback_shape_inference")
self.allow_onnxruntime = args_util.get(args, "allow_onnxruntime")
# No point is running ONNX shape inference if we're going to use fallback inference.
if self.force_fallback:
self.do_shape_inference = False
[docs] def add_to_script_impl(self, script, loader_name):
"""
Note that this method does not take fallback shape inference into account.
To support fallback shape inference, the tool must call `fallback_inference()` manually.
Args:
loader_name (str):
The name of the loader which should be consumed by the ``InferShapes`` loader.
Returns:
str: The name of the ``InferShapes`` loader added to the script.
"""
if self.do_shape_inference:
script.add_import(imports=["InferShapes"], frm="polygraphy.backend.onnx")
loader_name = script.add_loader(
make_invocable(
"InferShapes",
loader_name,
external_data_dir=self.arg_groups[OnnxLoadArgs].external_data_dir,
allow_onnxruntime=self.allow_onnxruntime,
),
"infer_shapes",
)
return loader_name
[docs] def infer_shapes(self, model, force=None):
"""
Run shape inference on an ONNX model if `do_shape_inference` is True
according to arguments provided on the command-line.
Args:
model (onnx.ModelProto): The model in which to infer shapes.
force (bool):
Force shape inference to run even if `do_shape_inference` is False.
Defaults to False.
Returns:
onnx.ModelProto: The model with shapes inferred.
"""
force = util.default(force, False)
with util.TempAttrChange(self, {"do_shape_inference": True if force else self.do_shape_inference}):
loader = args_util.run_script(self.add_to_script, model)
return util.invoke_if_callable(loader)[0]
[docs] def fallback_inference(self, onnx_model, outputs=None):
"""
Run inference with ONNX-Runtime.
This can be used to retrieve values/shapes/data types for all
tensors in the model when other shape inference approaches fail.
Args:
onnx_model (onnx.ModelProto):
The ONNX model in which to infer shapes.
outputs (List[str]):
The names of the outputs to retrieved.
Defaults to constants.MARK_ALL
Returns:
(IterationResult, TensorMetadata):
A tuple containing two elements:
1. Mapping of values for all tensors in the model, including inputs.
2. Metadata for every tensor in the model.
"""
outputs = util.default(outputs, constants.MARK_ALL)
with G_LOGGER.verbosity(G_LOGGER.module_severity.get(G_LOGGER.module_path(__file__)) + 10):
load_model = onnx_backend.ModifyOutputs(onnx_model, outputs=outputs, copy=True)
with onnxrt_backend.OnnxrtRunner(
onnxrt_backend.SessionFromOnnx(onnx_backend.BytesFromOnnx(load_model))
) as runner:
data_loader = self.arg_groups[DataLoaderArgs].get_data_loader()
loader_cache = DataLoaderCache(data_loader)
loader_cache.set_input_metadata(runner.get_input_metadata())
feed_dict = loader_cache[0]
with G_LOGGER.verbosity(G_LOGGER.module_severity.get(G_LOGGER.module_path(__file__)) - 10):
G_LOGGER.info(
f"Running fallback shape inference using input metadata:\n{TensorMetadata.from_feed_dict(feed_dict)}"
)
outputs = runner.infer(feed_dict)
# We include the inputs here so that we have values for all tensors in the model.
outputs.update(feed_dict)
# Use IterationResult here since it can handle very large tensors by saving to disk.
# Layerwise outputs might otherwise take up too much memory.
return IterationResult(outputs), TensorMetadata.from_feed_dict(outputs)
[docs]@mod.export()
class OnnxSaveArgs(BaseArgs):
"""
ONNX Model Saving: saving ONNX models.
Depends on:
- OnnxInferShapesArgs: if allow_shape_inference == True
"""
def __init__(
self,
allow_shape_inference: bool = None,
output_opt: str = None,
output_short_opt: str = None,
output_opt_required: bool = None,
output_default_path: str = None,
allow_multiple_models: bool = None,
):
"""
Args:
allow_shape_inference (bool):
Whether to allow shape inference when saving models.
Defaults to False.
output_opt (str):
The name of the output path option.
Defaults to "output".
Use a value of ``False`` to disable the option.
output_short_opt (str):
The short option to use for the output path.
Defaults to "-o".
Use a value of ``False`` to disable the short option.
output_opt_required (bool):
Whether the output path is a required argument.
Defaults to False.
output_default_path (str):
The default value to use for the output path option.
Defaults to None.
allow_multiple_models (bool):
Whether to enable support for saving more than one model.
If this is True, the output path is expected to be a directory.
Defaults to False.
"""
super().__init__()
self._allow_shape_inference = util.default(allow_shape_inference, False)
self._output_opt = util.default(output_opt, "output")
self._output_short_opt = util.default(output_short_opt, "-o")
self._output_opt_required = util.default(output_opt_required, False)
self._output_default_path = output_default_path
self._allow_multiple_models = util.default(allow_multiple_models, False)
# add_to_script should never be called when `allow_multiple_models` is enabled.
# The one exception is that `save_onnx` should be able to call it, which is why we need this escape hatch.
self._disable_add_to_script_check = False
def add_parser_args_impl(self):
if self._output_opt:
params = ([self._output_short_opt] if self._output_short_opt else []) + [f"--{self._output_opt}"]
help_msg = "Path to save the ONNX model"
if self._allow_multiple_models:
help_msg = "Path to a directory in which to save ONNX model(s)"
self.group.add_argument(
*params,
help=help_msg,
dest="save_onnx",
default=self._output_default_path,
required=self._output_opt_required,
)
self.group.add_argument(
"--save-external-data",
"--external-data-path",
help="Whether to save weight data in external file(s). "
+ (
"You may optionally provide a value to this argument which will be used as a suffix for the external data files"
if self._allow_multiple_models
else "To use a non-default path, supply the desired path as an argument. This is always a relative path; "
"external data is always written to the same directory as the model. "
),
default=None,
action="append",
nargs="?",
dest="external_data_path",
)
self.group.add_argument(
"--external-data-size-threshold",
help="The size threshold, in bytes, above which tensor data will be stored in the external file. "
"Tensors smaller that this threshold will remain in the ONNX file. "
"Optionally, use a `K`, `M`, or `G` suffix to indicate KiB, MiB, or GiB respectively. "
"For example, `--external-data-size-threshold=16M` is equivalent to `--external-data-size-threshold=16777216`. "
"Has no effect if `--save-external-data` is not set. Defaults to 1024 bytes.",
default=None,
)
self.group.add_argument(
"--no-save-all-tensors-to-one-file",
help="Do not save all tensors to a single file when saving external data. "
"Has no effect if `--save-external-data` is not set",
dest="all_tensors_to_one_file",
default=None,
action="store_false",
)
[docs] def parse_impl(self, args):
"""
Parses command-line arguments and populates the following attributes:
Attributes:
path (str): The path at which to save the ONNX model.
external_data_path (str): The path at which to save external data.
size_threshold (int): The size threshold above which external data is saved.
all_tensors_to_one_file (bool): Whether all external data should be written to a single file.
"""
self.path = args_util.get(args, "save_onnx")
external_data_path = args_util.get(args, "external_data_path")
if external_data_path is not None:
external_data_path = external_data_path[0] or ""
self.external_data_path = external_data_path
self.size_threshold = args_util.parse_num_bytes(args_util.get(args, "external_data_size_threshold"))
self.all_tensors_to_one_file = args_util.get(args, "all_tensors_to_one_file")
[docs] def add_to_script_impl(self, script, loader_name):
"""
Args:
loader_name (str):
The name of the loader which should be consumed by the ``SaveOnnx`` loader.
Returns:
str: The name of the ``SaveOnnx`` loader added to the script.
"""
if self._allow_multiple_models and not self._disable_add_to_script_check:
G_LOGGER.internal_error(
"OnnxSaveArgs.add_to_script() should never be called when `allow_multiple_models` is enabled"
)
if self.path is None:
return loader_name
# Need to run shape inference again after processing the graph since it may have changed.
if self._allow_shape_inference:
loader_name = self.arg_groups[OnnxInferShapesArgs].add_to_script(script, loader_name)
script.add_import(imports=["SaveOnnx"], frm="polygraphy.backend.onnx")
loader_name = script.add_loader(
make_invocable(
"SaveOnnx",
loader_name,
path=self.path,
external_data_path=self.external_data_path,
size_threshold=self.size_threshold,
all_tensors_to_one_file=self.all_tensors_to_one_file,
),
"save_onnx",
)
return loader_name
[docs] def save_onnx(self, model, path: str = None):
"""
Saves an ONNX model according to arguments provided on the command-line.
Args:
model (onnx.ModelProto): The ONNX model to save.
path (str):
The path at which to save the model.
If no path is provided, it is determined from command-line arguments.
Returns:
onnx.ModelProto: The model that was saved.
"""
attrs = {"path": path, "_disable_add_to_script_check": True}
if self._allow_multiple_models:
if self.external_data_path is not None:
attrs["external_data_path"] = os.path.basename(os.path.splitext(path)[0]) + (
self.external_data_path or "_ext_data"
)
with util.TempAttrChange(self, attrs):
loader = args_util.run_script(self.add_to_script, model)
return loader()
[docs]@mod.export()
class OnnxLoadArgs(BaseArgs):
"""
ONNX Model Loading: loading ONNX models.
Depends on:
- ModelArgs
- OnnxInferShapesArgs: if allow_shape_inference == True
- OnnxSaveArgs: if allow_saving == True
- OnnxFromTfArgs: if allow_from_tf == True
"""
def __init__(
self,
allow_saving: bool = None,
outputs_opt_prefix: str = None,
allow_shape_inference: bool = None,
allow_from_tf: bool = None,
):
"""
Args:
allow_saving (bool):
Whether to allow loaded models to be saved.
Defaults to False.
outputs_opt_prefix (str):
The prefix to use for the outputs option, which controls which tensors are marked as outputs.
Defaults to "onnx-".
Use a value of ``False`` to disable the option.
allow_shape_inference (bool):
Whether to allow shape inference when loading models.
Defaults to True.
allow_from_tf (bool):
Whether to allow conversion of TensorFlow models to ONNX.
Defaults to False.
"""
super().__init__()
self._allow_saving = util.default(allow_saving, False)
self._allow_shape_inference = util.default(allow_shape_inference, True)
self._outputs_opt_prefix = util.default(outputs_opt_prefix, "onnx-")
self._allow_from_tf = util.default(allow_from_tf, False)
def add_parser_args_impl(self):
self.group.add_argument(
"--external-data-dir",
"--load-external-data",
"--ext",
dest="external_data_dir",
help="Path to a directory containing external data for the model. "
"Generally, this is only required if the external data is not stored in the model directory.",
)
self.group.add_argument(
"--ignore-external-data",
help="Ignore external data and just load the model structure without any weights. "
"The model will be usable only for purposes that don't require weights, such as extracting "
"subgraphs or inspecting model structure. "
"This can be useful in cases where external data is not available.",
action="store_true",
default=None,
)
if self._outputs_opt_prefix is not False: # Empty strings should not disable the option
self.group.add_argument(
f"--{self._outputs_opt_prefix}outputs",
help="Name(s) of ONNX tensor(s) to mark as output(s). "
"Using the special value 'mark all' indicates that all tensors should be used as outputs",
nargs="+",
default=None,
dest="onnx_outputs",
)
self.group.add_argument(
f"--{self._outputs_opt_prefix}exclude-outputs",
help="[EXPERIMENTAL] Name(s) of ONNX output(s) to unmark as outputs.",
nargs="+",
default=None,
dest="onnx_exclude_outputs",
)
[docs] def parse_impl(self, args):
"""
Parses command-line arguments and populates the following attributes:
Attributes:
outputs (List[str]): Names of output tensors.
exclude_outputs (List[str]): Names of tensors which should be unmarked as outputs.
external_data_dir (str): Path to a directory from which to load external data.
ignore_external_data (bool): Whether to ignore loading external data.
"""
self.outputs = args_util.get_outputs(args, "onnx_outputs")
self.exclude_outputs = args_util.get(args, "onnx_exclude_outputs")
self.external_data_dir = args_util.get(args, "external_data_dir")
self.ignore_external_data = args_util.get(args, "ignore_external_data")
def _add_modify_onnx_outputs(self, script, loader_name, disable_custom_outputs: bool = None):
if disable_custom_outputs:
outputs = None
exclude_outputs = None
else:
outputs = args_util.get_outputs_for_script(script, self.outputs)
exclude_outputs = self.exclude_outputs
modify_outputs_loader = make_invocable_if_nondefault_kwargs(
"ModifyOnnxOutputs", loader_name, outputs=outputs, exclude_outputs=exclude_outputs
)
if modify_outputs_loader is not None:
script.add_import(imports="ModifyOutputs", frm="polygraphy.backend.onnx", imp_as="ModifyOnnxOutputs")
loader_name = script.add_loader(
modify_outputs_loader,
"modify_outputs",
)
return loader_name
[docs] def add_to_script_impl(self, script, disable_custom_outputs: bool = None, serialize_model: bool = None):
"""
Args:
disable_custom_outputs (bool):
Whether to disallow modifying outputs according to the `outputs` and `exclude_outputs` attributes.
Defaults to False.
serialize_model (bool):
Whether to serialize the model.
Defaults to False.
Returns:
str: The name of the ONNX loader added in the script.
"""
model_type = self.arg_groups[ModelArgs].model_type
if model_type.is_onnx():
loader_name = self.arg_groups[ModelArgs].path
if self._allow_shape_inference:
loader_name = self.arg_groups[OnnxInferShapesArgs].add_to_script(script, loader_name)
if loader_name == self.arg_groups[ModelArgs].path: # Shape inference loader isn't being used, have to load.
script.add_import(imports=["OnnxFromPath"], frm="polygraphy.backend.onnx")
loader_str = make_invocable(
"OnnxFromPath",
self.arg_groups[ModelArgs].path,
external_data_dir=self.external_data_dir,
ignore_external_data=self.ignore_external_data,
)
loader_name = script.add_loader(loader_str, "load_onnx")
elif model_type.is_tf() and self._allow_from_tf:
from polygraphy.tools.args.backend.onnx.loader import OnnxFromTfArgs
loader_name = self.arg_groups[OnnxFromTfArgs].add_to_script(script)
else:
G_LOGGER.critical(f"Model type: {model_type} could not be converted to an ONNX model.")
loader_name = self._add_modify_onnx_outputs(script, loader_name, disable_custom_outputs=disable_custom_outputs)
if self._allow_saving:
loader_name = self.arg_groups[OnnxSaveArgs].add_to_script(script, loader_name)
if serialize_model:
script.add_import(imports=["BytesFromOnnx"], frm="polygraphy.backend.onnx")
loader_name = script.add_loader(make_invocable("BytesFromOnnx", loader_name), "serialize_onnx")
return loader_name
[docs] def must_use_onnx_loader(self, disable_custom_outputs: bool = None):
"""
Whether this model needs to be loaded via a Polygraphy ONNX loader, e.g., in case it
needs modifications.
Args:
disable_custom_outputs (bool):
Whether to disallow modifying outputs according to the `outputs` and `exclude_outputs` attributes.
Returns:
bool
"""
tmp_script = Script()
inp_loader = "check_needs_modify"
needs_modify = self._add_modify_onnx_outputs(tmp_script, inp_loader, disable_custom_outputs) != inp_loader
needs_shape_inference = self._allow_shape_inference and self.arg_groups[OnnxInferShapesArgs].do_shape_inference
needs_save = self._allow_saving and self.arg_groups[OnnxSaveArgs].path is not None
# Currently, other loaders do not support external data, so we must fall back to the ONNX loader if it's present.
return (
not self.arg_groups[ModelArgs].model_type.is_onnx()
or needs_modify
or self.external_data_dir
or needs_shape_inference
or needs_save
)
[docs] def load_onnx(self):
"""
Loads an ONNX model according to arguments provided on the command-line.
Returns:
onnx.ModelProto: The model that was loaded.
"""
loader = args_util.run_script(self.add_to_script)
return loader()
[docs]@mod.export()
class OnnxFromTfArgs(BaseArgs):
"""
TensorFlow-ONNX Model Conversion: converting TensorFlow models to ONNX.
Depends on:
- TfLoadArgs
"""
def add_parser_args_impl(self):
self.group.add_argument("--opset", help="Opset to use when converting to ONNX", default=None, type=int)
[docs] def parse_impl(self, args):
"""
Parses command-line arguments and populates the following attributes:
Attributes:
opset (int): The ONNX opset version to use during conversion.
"""
self.opset = args_util.get(args, "opset")
def add_to_script_impl(self, script):
from polygraphy.tools.args.backend.tf.loader import TfLoadArgs
G_LOGGER.verbose(
"Attempting to load as a TensorFlow model, using TF2ONNX to convert to ONNX. "
"If this is not correct, please specify --model-type",
mode=LogMode.ONCE,
)
script.add_import(imports=["OnnxFromTfGraph"], frm="polygraphy.backend.onnx")
loader_str = make_invocable(
"OnnxFromTfGraph",
self.arg_groups[TfLoadArgs].add_to_script(script, disable_custom_outputs=True),
opset=self.opset,
)
loader_name = script.add_loader(loader_str, "export_onnx_from_tf")
return loader_name