def export_onnx(
self,
name: str = "model.onnx",
opset: int = DEFAULT_ONNX_OPSET,
doc_string: str = "",
debug_mode: bool = True,
**kwargs,
):
"""
Export an ONNX file for the current model.
:param name: name of the onnx file to save
:param opset: onnx opset to use for exported model. Default is 11
:param doc_string: optional doc string for exported ONNX model
:param debug_mode: debug mode, default to True, passed into `convert_keras`
:param kwargs: additional parameters passed into `convert_keras`
"""
if keras2onnx_import_error is not None:
raise keras2onnx_import_error
model_name = self._model.name or name.split(".onnx")[0]
onnx_model = keras2onnx.convert_keras(
self._model,
name=model_name,
target_opset=opset,
doc_string=doc_string,
debug_mode=debug_mode,
**kwargs,
)
onnx_path = os.path.join(self._output_dir, name)
create_parent_dirs(onnx_path)
keras2onnx.save_model(onnx_model, onnx_path)
def save_framework_info(framework: Any, path: Optional[str] = None):
"""
Save the framework info for a given framework.
If path is provided, will save to a json file at that path.
If path is not provided, will print out the info.
:param framework: The item to detect the ML framework for.
See :func:`detect_framework` for more information.
:type framework: Any
:param path: The path, if any, to save the info to in json format.
If not provided will print out the info.
:type path: Optional[str]
"""
_LOGGER.debug(
"saving framework info for framework %s to %s",
framework,
path if path else "sys.out",
)
info = (framework_info(framework)
if not isinstance(framework, FrameworkInfo) else framework)
if path:
path = clean_path(path)
create_parent_dirs(path)
with open(path, "w") as file:
file.write(info.json())
_LOGGER.info("saved framework info for framework %s in file at %s",
framework, path),
else:
print(info.json(indent=4))
_LOGGER.info("printed out framework info for framework %s", framework)
def export_pytorch(
self,
optimizer: Optimizer = None,
epoch: int = None,
name: str = "model.pth",
use_zipfile_serialization_if_available: bool = True,
include_modifiers: bool = False,
):
"""
Export the pytorch state dicts into pth file within a
pytorch framework directory.
:param optimizer: optional optimizer to export along with the module
:param epoch: optional epoch to export along with the module
:param name: name of the pytorch file to save
:param use_zipfile_serialization_if_available: for torch >= 1.6.0 only
exports the Module's state dict using the new zipfile serialization
:param include_modifiers: if True, and a ScheduledOptimizer is provided
as the optimizer, the associated ScheduledModifierManager and its
Modifiers will be exported under the 'manager' key. Default is False
"""
pytorch_path = os.path.join(self._output_dir, "pytorch")
pth_path = os.path.join(pytorch_path, name)
create_parent_dirs(pth_path)
save_model(
pth_path,
self._module,
optimizer,
epoch,
use_zipfile_serialization_if_available=(
use_zipfile_serialization_if_available),
include_modifiers=include_modifiers,
)
def export_pb(
self,
outputs: List[Union[str, tf_compat.Tensor]],
graph: tf_compat.Graph = None,
sess: tf_compat.Session = None,
):
"""
Export a serialized pb version of the a graph and session.
:param outputs: the list of outputs the graph should be created for
(used to determine the scope of the graph to export),
can be either a list of names or a list of tensors
:param graph: the graph to export to a pb format,
if not supplied will use get_default_graph()
:param sess: the session to export to a pb format,
if not supplied will use get_default_session()
"""
if not graph:
graph = tf_compat.get_default_graph()
if not sess:
sess = tf_compat.get_default_session()
outputs = [
out if isinstance(out, str) else clean_tensor_name(out) for out in outputs
]
output_graph_def = tf_compat.graph_util.convert_variables_to_constants(
sess, graph.as_graph_def(), outputs
)
create_parent_dirs(self.pb_path)
with tf_compat.gfile.GFile(self.pb_path, "wb") as f:
f.write(output_graph_def.SerializeToString())
def save_benchmark_results(
model: Any,
data: Any,
batch_size: int,
iterations: int,
warmup_iterations: int,
framework: Optional[str],
provider: Optional[str] = None,
device: Optional[str] = None,
save_path: Optional[str] = None,
framework_args: Dict[str, Any] = {},
show_progress: bool = False,
):
"""
Saves the benchmark results ran for specific framework.
If path is provided, will save to a json file at the path.
If path is not provided, will print out the info.
If no framework is provided, will detect the framework based on the model.
:param model: model to benchmark
:param data: data to benchmark
:param batch_size: batch size
:param iterations: number of iterations
:param warmup_iterations: number of warmup iterations
:param framework: the specific framework run the benchmark in
:param provider: the specific inference provider to use
:param device: the specific device to use
:param save_path: path to save the benchmark results
:param framework_args: additional framework specific arguments to
pass to the runner
:param show_progress: True to show a tqdm bar when running, False otherwise
"""
results = execute_in_sparseml_framework(
framework if framework is not None else model,
"run_benchmark",
model,
data,
batch_size=batch_size,
iterations=iterations,
warmup_iterations=warmup_iterations,
provider=provider,
device=device,
framework_args=framework_args,
show_progress=show_progress,
)
if save_path:
save_path = clean_path(save_path)
create_parent_dirs(save_path)
with open(save_path, "w") as file:
file.write(results.json(indent=4))
_LOGGER.info(f"saved benchmark results in file at {save_path}"),
else:
print(results.json(indent=4))
_LOGGER.info("printed out benchmark results")
def save(self, file_path: str):
"""
:param file_path: the file path to save the yaml config representation to
"""
file_path = clean_path(file_path)
create_parent_dirs(file_path)
with open(file_path, "w") as yaml_file:
yaml_file.write(str(self))
def save(self, file_path: str):
"""
saves the dict representation of this ModelInfo object as a JSON file
to the given file path
:param file_path: file path to save to
"""
create_parent_dirs(file_path)
with open(file_path, "w") as file:
json.dump(self.to_dict(), file)
def save_descs(descs: List, path: str):
"""
Save a list of AnalyzedLayerDesc to a json file
:param descs: a list of descriptions to save
:param path: the path to save the descriptions at
"""
path = clean_path(path)
create_parent_dirs(path)
save_obj = {"descriptions": [desc.dict() for desc in descs]}
with open(path, "w") as file:
json.dump(save_obj, file)
def save_json(self, path: str):
"""
:param path: the path to save the json file at representing the analyzed
results
"""
if not path.endswith(".json"):
path += ".json"
path = clean_path(path)
create_parent_dirs(path)
with open(path, "w") as file:
dictionary = self.dict()
json.dump(dictionary, file, indent=2)
def export_onnx(
self,
name: str = "model.onnx",
opset: int = DEFAULT_ONNX_OPSET,
doc_string: str = "",
debug_mode: bool = True,
raise_on_tf_support: bool = True,
**kwargs,
):
"""
Export an ONNX file for the current model.
:param name: name of the onnx file to save
:param opset: onnx opset to use for exported model. Default is 11
:param doc_string: optional doc string for exported ONNX model
:param debug_mode: debug mode, default to True, passed into `convert_keras`
:param kwargs: additional parameters passed into `convert_keras`
"""
if keras2onnx_import_error is not None:
raise keras2onnx_import_error
if raise_on_tf_support:
import tensorflow
v = tensorflow.__version__
if v >= "2.3.0":
raise ValueError(
f"Tensorflow version {v} is greater than the currently supported "
"version for keras2onnx. Please downgrade the Tensorflow <2.3.0 "
"or set raise_on_tf_support to False to continue.")
model_name = self._model.name or name.split(".onnx")[0]
onnx_model = keras2onnx.convert_keras(
self._model,
name=model_name,
target_opset=opset,
doc_string=doc_string,
debug_mode=debug_mode,
**kwargs,
)
onnx_path = os.path.join(self._output_dir, name)
create_parent_dirs(onnx_path)
keras2onnx.save_model(onnx_model, onnx_path)
def export_torchscript(
self,
name: str = "model.pts",
sample_batch: Optional[Any] = None,
):
"""
Export the torchscript into a pts file within a framework directory. If
a sample batch is provided, will create torchscript model in trace mode.
Otherwise uses script to create torchscript.
:param name: name of the torchscript file to save
:param sample_batch: If provided, will create torchscript model via tracing
using the sample_batch
"""
path = os.path.join(self._output_dir, "framework", name)
create_parent_dirs(path)
if sample_batch:
trace_model(path, self._module, sample_batch)
else:
script_model(path, self._module)
def export_checkpoint(
self, saver: tf_compat.train.Saver = None, sess: tf_compat.Session = None
):
"""
Export a checkpoint for the current TensorFlow graph and session.
:param saver: the saver instance to use to save the current session,
if not supplied will create a new one using TRAINABLE_VARIABLES
:param sess: the current session to export a checkpoint for,
if not supplied will use get_default_session()
"""
if not sess:
sess = tf_compat.get_default_session()
if not saver:
saver = tf_compat.train.Saver(
tf_compat.get_collection(tf_compat.GraphKeys.TRAINABLE_VARIABLES)
)
create_parent_dirs(self.checkpoint_path)
saver.save(sess, self.checkpoint_path)
def _save_keras(self, file_path: str):
create_parent_dirs(file_path)
self._model.save(file_path)
def export_onnx(
self,
sample_batch: Any,
name: str = "model.onnx",
opset: int = DEFAULT_ONNX_OPSET,
disable_bn_fusing: bool = True,
):
"""
Export an onnx file for the current module and for a sample batch.
Sample batch used to feed through the model to freeze the graph for a
particular execution.
:param sample_batch: the batch to export an onnx for, handles creating the
static graph for onnx as well as setting dimensions
:param name: name of the onnx file to save
:param opset: onnx opset to use for exported model. Default is 11, if torch
version is 1.2 or below, default is 9
:param disable_bn_fusing: torch >= 1.7.0 only. Set True to disable batch norm
fusing during torch export. Default and suggested setting is True. Batch
norm fusing will change the exported parameter names as well as affect
sensitivity analyses of the exported graph. Additionally, the DeepSparse
inference engine, and other engines, perform batch norm fusing at model
compilation.
"""
sample_batch = tensors_to_device(sample_batch, "cpu")
onnx_path = os.path.join(self._output_dir, name)
create_parent_dirs(onnx_path)
with torch.no_grad():
out = tensors_module_forward(sample_batch, self._module)
input_names = None
if isinstance(sample_batch, Tensor):
input_names = ["input"]
elif isinstance(sample_batch, Iterable):
input_names = [
"input_{}".format(index)
for index, _ in enumerate(iter(sample_batch))
]
output_names = None
if isinstance(out, Tensor):
output_names = ["output"]
elif isinstance(out, Iterable):
output_names = [
"output_{}".format(index) for index, _ in enumerate(iter(out))
]
# disable active quantization observers because they cannot be exported
disabled_observers = []
for submodule in self._module.modules():
if (hasattr(submodule, "observer_enabled")
and submodule.observer_enabled[0] == 1):
submodule.observer_enabled[0] = 0
disabled_observers.append(submodule)
is_quant_module = any(
hasattr(submodule, "qconfig") and submodule.qconfig
for submodule in self._module.modules())
batch_norms_wrapped = False
if torch.__version__ >= "1.7" and not is_quant_module and disable_bn_fusing:
# prevent batch norm fusing by adding a trivial operation before every
# batch norm layer
export_module = deepcopy(self._module)
batch_norms_wrapped = _wrap_batch_norms(export_module)
else:
export_module = self._module
torch.onnx.export(
export_module,
sample_batch,
onnx_path,
input_names=input_names,
output_names=output_names,
strip_doc_string=True,
verbose=False,
opset_version=opset,
)
# re-enable disabled quantization observers
for submodule in disabled_observers:
submodule.observer_enabled[0] = 1
# clean up graph from any injected / wrapped operations
if batch_norms_wrapped:
onnx_model = onnx.load(onnx_path)
_delete_trivial_onnx_adds(onnx_model)
onnx.save(onnx_model, onnx_path)
def pb_to_onnx(
inputs: List[Union[str, tf_compat.Tensor]],
outputs: List[Union[str, tf_compat.Tensor]],
pb_path: str,
onnx_path: str,
opset: int = default_onnx_opset(),
custom_op_handlers=None,
extra_opset=None,
shape_override: Dict[str, List] = None,
):
"""
Export an ONNX format for the graph from PB format.
Should not be called within an active graph or session.
:param inputs: the inputs the graph should be created for,
can be either a list of names or a list of tensors
:param outputs: the outputs the graph should be created for,
can be either a list of names or a list of tensors
:param pb_path: path to the existing PB file
:param onnx_path: path to the output ONNX file
:param opset: ONNX opset
:param custom_op_handlers: dictionary of custom op handlers
:param extra_opset: list of extra opset's
:param shape_override: new shape to override
"""
try:
from tf2onnx import constants, optimizer, utils
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
except ModuleNotFoundError:
raise ModuleNotFoundError(
"tf2onnx must be installed on the system before using export_onnx"
)
try:
from tf2onnx import tf_loader as loader
except Exception:
from tf2onnx import loader
pb_path = clean_path(pb_path)
if not os.path.exists(pb_path):
raise FileNotFoundError(
("no pb file for the model found at {}").format(pb_path)
)
inputs = [inp if isinstance(inp, str) else inp.name for inp in inputs]
outputs = [out if isinstance(out, str) else out.name for out in outputs]
graph_def, inputs, outputs = loader.from_graphdef(pb_path, inputs, outputs)
graph_def = tf_optimize(inputs, outputs, graph_def, fold_constant=True)
with tf_compat.Graph().as_default() as tf_graph:
tf_compat.import_graph_def(graph_def, name="")
with tf_compat.Session(graph=tf_graph):
graph = process_tf_graph(
tf_graph,
continue_on_error=False,
target=",".join(constants.DEFAULT_TARGET),
opset=opset,
custom_op_handlers=custom_op_handlers,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=inputs,
output_names=outputs,
)
onnx_graph = optimizer.optimize_graph(graph)
model_proto = onnx_graph.make_model("converted from {}".format(pb_path))
onnx_path = clean_path(onnx_path)
create_parent_dirs(onnx_path)
utils.save_protobuf(onnx_path, model_proto)
def _save_file_str(content: str, file_path: str):
create_parent_dirs(file_path)
with open(file_path, "w") as file:
file.write(content)
def export_onnx(
module: Module,
sample_batch: Any,
file_path: str,
opset: int = DEFAULT_ONNX_OPSET,
disable_bn_fusing: bool = True,
convert_qat: bool = False,
dynamic_axes: Union[str, Dict[str, List[int]]] = None,
skip_input_quantize: bool = False,
**export_kwargs,
):
"""
Export an onnx file for the current module and for a sample batch.
Sample batch used to feed through the model to freeze the graph for a
particular execution.
:param module: torch Module object to export
:param sample_batch: the batch to export an onnx for, handles creating the
static graph for onnx as well as setting dimensions
:param file_path: path to the onnx file to save
:param opset: onnx opset to use for exported model. Default is 11, if torch
version is 1.2 or below, default is 9
:param disable_bn_fusing: torch >= 1.7.0 only. Set True to disable batch norm
fusing during torch export. Default and suggested setting is True. Batch
norm fusing will change the exported parameter names as well as affect
sensitivity analyses of the exported graph. Additionally, the DeepSparse
inference engine, and other engines, perform batch norm fusing at model
compilation.
:param convert_qat: if True and quantization aware training is detected in
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 dynamic_axes: dictionary of input or output names to list of dimensions
of those tensors that should be exported as dynamic. May input 'batch'
to set the first dimension of all inputs and outputs to dynamic. Default
is an empty dict
:param skip_input_quantize: if True, the export flow will attempt to delete
the first Quantize Linear Nodes(s) immediately after model input and set
the model input type to UINT8. 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):
warnings.warn(
"Sample inputs passed into the ONNX exporter should be in "
"the same order defined in the model forward function. "
"Consider using OrderedDict for this purpose.",
UserWarning,
)
sample_batch = tensors_to_device(sample_batch, "cpu")
create_parent_dirs(file_path)
module = deepcopy(module).cpu()
module.eval()
with torch.no_grad():
out = tensors_module_forward(sample_batch,
module,
check_feat_lab_inp=False)
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
if "output_names" not in export_kwargs:
export_kwargs["output_names"] = _get_output_names(out)
if dynamic_axes == "batch":
dynamic_axes = {
tensor_name: {
0: "batch"
}
for tensor_name in (export_kwargs["input_names"] +
export_kwargs["output_names"])
}
# disable active quantization observers because they cannot be exported
disabled_observers = []
for submodule in module.modules():
if (hasattr(submodule, "observer_enabled")
and submodule.observer_enabled[0] == 1):
submodule.observer_enabled[0] = 0
disabled_observers.append(submodule)
is_quant_module = any(
hasattr(submodule, "qconfig") and submodule.qconfig
for submodule in module.modules())
batch_norms_wrapped = False
if torch.__version__ >= "1.7" and not is_quant_module and disable_bn_fusing:
# prevent batch norm fusing by adding a trivial operation before every
# batch norm layer
batch_norms_wrapped = _wrap_batch_norms(module)
torch.onnx.export(
module,
sample_batch,
file_path,
strip_doc_string=True,
verbose=False,
opset_version=opset,
dynamic_axes=dynamic_axes,
**export_kwargs,
)
# re-enable disabled quantization observers
for submodule in disabled_observers:
submodule.observer_enabled[0] = 1
# onnx file fixes
onnx_model = onnx.load(file_path)
# fix changed batch norm names
_fix_batch_norm_names(onnx_model)
if batch_norms_wrapped:
# clean up graph from any injected / wrapped operations
_delete_trivial_onnx_adds(onnx_model)
onnx.save(onnx_model, file_path)
if convert_qat and is_quant_module:
# overwrite exported model with fully quantized version
quantize_torch_qat_export(model=file_path, output_file_path=file_path)
if skip_input_quantize:
try:
skip_onnx_input_quantize(file_path, file_path)
except Exception as e:
_LOGGER.warning(
f"Unable to skip input QuantizeLinear op with exception {e}")