def save_model_to_file(self, output_path, use_external_data_format=False, all_tensors_to_one_file=True):
logger.info(f"Sort graphs in topological order")
self.topological_sort()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
if output_path.endswith(".json"): # Output text for testing small model.
assert isinstance(self.model, ModelProto)
with open(output_path, "w") as out:
out.write(str(self.model))
else:
# Save model to external data, which is needed for model size > 2GB
if use_external_data_format:
output_dir = Path(output_path).parent
output_dir.mkdir(parents=True, exist_ok=True)
location = Path(output_path).name + ".data" if all_tensors_to_one_file else None
# Show warnings of potential confliction of existing external data file.
if all_tensors_to_one_file:
if os.path.exists(location):
logger.warning(
f"External data file ({location}) existed. Please remove the file and try again.")
else:
if os.listdir(output_dir):
logger.warning(
f"Output directory ({output_dir}) for external data is not empty. Please try again with a new directory."
)
external_data_helper.convert_model_to_external_data(self.model,
all_tensors_to_one_file=all_tensors_to_one_file,
location=location)
save_model(self.model, output_path)
logger.info(f"Model saved to {output_path}")
def test_convert_model_to_external_data_from_one_file_with_location(
self): # type: () -> None
model_file_path = self.get_temp_model_filename()
external_data_file = str(uuid.uuid4())
convert_model_to_external_data(self.model,
size_threshold=0,
all_tensors_to_one_file=True,
location=external_data_file)
onnx.save_model(self.model, model_file_path)
self.assertTrue(
Path.isfile(os.path.join(self.temp_dir, external_data_file)))
model = onnx.load_model(model_file_path)
# test convert model from external data
convert_model_from_external_data(model)
model_file_path = self.get_temp_model_filename()
onnx.save_model(model, model_file_path)
model = onnx.load_model(model_file_path)
initializer_tensor = model.graph.initializer[0]
self.assertFalse(len(initializer_tensor.external_data))
self.assertEqual(initializer_tensor.data_location, TensorProto.DEFAULT)
self.assertTrue(
np.allclose(to_array(initializer_tensor), self.initializer_value))
attribute_tensor = model.graph.node[0].attribute[0].t
self.assertFalse(len(attribute_tensor.external_data))
self.assertEqual(attribute_tensor.data_location, TensorProto.DEFAULT)
self.assertTrue(
np.allclose(to_array(attribute_tensor), self.attribute_value))
def test_convert_model_to_external_data_from_one_file_without_location_uses_model_name(self): # type: () -> None
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model, size_threshold=0, all_tensors_to_one_file=True)
onnx.save_model(self.model, model_file_path)
self.assertTrue(Path.isfile(model_file_path))
self.assertTrue(Path.isfile(os.path.join(self.temp_dir, model_file_path)))
def save_model(proto, f, format=None, save_as_external_data=False, all_tensors_to_one_file=True, location=None, size_threshold=1024, convert_attribute=False):
if isinstance(proto, bytes):
proto = onnx._deserialize(proto, onnx.ModelProto())
if save_as_external_data:
convert_model_to_external_data(proto, all_tensors_to_one_file, location, size_threshold, convert_attribute)
s = onnx._serialize(proto)
onnx._save_bytes(s, f)
def test_convert_model_to_external_data_with_size_threshold(self) -> None:
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model, size_threshold=1024)
onnx.save_model(self.model, model_file_path)
model = onnx.load_model(model_file_path)
initializer_tensor = model.graph.initializer[0]
self.assertFalse(initializer_tensor.HasField("data_location"))
def test_convert_model_to_external_data_one_file_per_tensor_with_attribute(self): # type: () -> None
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model, size_threshold=0, all_tensors_to_one_file=False, convert_attribute=True)
onnx.save_model(self.model, model_file_path)
self.assertTrue(Path.isfile(model_file_path))
self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "input_value")))
self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
def test_convert_model_to_external_data_without_size_threshold(self): # type: () -> None
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model, size_threshold=0)
onnx.save_model(self.model, model_file_path)
model = onnx.load_model(model_file_path)
initializer_tensor = model.graph.initializer[0]
self.assertTrue(initializer_tensor.HasField("data_location"))
self.assertTrue(np.allclose(to_array(initializer_tensor), self.initializer_value))
def test_convert_model_to_external_data_does_not_convert_attribute_values(self): # type: () -> None
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model, size_threshold=0, convert_attribute=False, all_tensors_to_one_file=False)
onnx.save_model(self.model, model_file_path)
self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "input_value")))
self.assertFalse(Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
model = onnx.load_model(model_file_path)
initializer_tensor = model.graph.initializer[0]
self.assertTrue(initializer_tensor.HasField("data_location"))
attribute_tensor = model.graph.node[0].attribute[0].t
self.assertFalse(attribute_tensor.HasField("data_location"))
def save_model_to_file(self, output_path, use_external_data_format=False):
logger.info(f"Output model to {output_path}")
if output_path.endswith(
".json"): # Output text for testing small model.
assert isinstance(self.model, ModelProto)
with open(output_path, "w") as out:
out.write(str(self.model))
else:
# Save model to external data, which is needed for model size > 2GB
if use_external_data_format:
from pathlib import Path
external_data_helper.convert_model_to_external_data(
self.model,
all_tensors_to_one_file=True,
location=Path(output_path).name + ".data")
save_model(self.model, output_path)
def save_model_to_file(self, output_path, use_external_data_format=False):
logger.info(f"Output model to {output_path}")
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
if output_path.endswith(".json"): # Output text for testing small model.
assert isinstance(self.model, ModelProto)
with open(output_path, "w") as out:
out.write(str(self.model))
else:
# Save model to external data, which is needed for model size > 2GB
if use_external_data_format:
data_file = str(Path(output_path).name + ".data")
if os.path.isfile(data_file):
os.remove(data_file)
external_data_helper.convert_model_to_external_data(self.model,
all_tensors_to_one_file=True,
location=data_file)
save_model(self.model, output_path)
def test_convert_model_to_external_data_one_file_per_tensor(
self): # type: () -> None
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model,
all_tensors_to_one_file=False)
onnx.save_model(self.model, model_file_path)
self.assertTrue(Path.isfile(model_file_path))
self.assertTrue(Path.isfile(os.path.join(self.temp_dir,
"input_value")))
self.assertTrue(
Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
model = onnx.load_model(model_file_path)
initializer_tensor = model.graph.initializer[0]
self.assertTrue(
np.allclose(to_array(initializer_tensor), self.initializer_value))
attribute_tensor = model.graph.node[0].attribute[0].t
self.assertTrue(
np.allclose(to_array(attribute_tensor), self.attribute_value))
def test_convert_model_to_external_data_converts_attribute_values(
self) -> None:
model_file_path = self.get_temp_model_filename()
convert_model_to_external_data(self.model,
size_threshold=0,
convert_attribute=True)
onnx.save_model(self.model, model_file_path)
model = onnx.load_model(model_file_path)
initializer_tensor = model.graph.initializer[0]
self.assertTrue(
np.allclose(to_array(initializer_tensor), self.initializer_value))
self.assertTrue(initializer_tensor.HasField("data_location"))
attribute_tensor = model.graph.node[0].attribute[0].t
self.assertTrue(
np.allclose(to_array(attribute_tensor), self.attribute_value))
self.assertTrue(attribute_tensor.HasField("data_location"))
def save_model(proto,
f,
format=None,
save_as_external_data=False,
all_tensors_to_one_file=True,
location=None,
size_threshold=1024,
convert_attribute=False):
# type: (Union[ModelProto, bytes], Union[IO[bytes], Text], Optional[Any], bool, bool, Optional[Text], int, bool) -> None
'''
Saves the ModelProto to the specified path and optionally, serialize tensors with raw data as external data before saving.
@params
proto: should be a in-memory ModelProto
f: can be a file-like object (has "write" function) or a string containing a file name format for future use
all_tensors_to_one_file: If true, save all tensors to one external file specified by location.
If false, save each tensor to a file named with the tensor name.
location: specify the external file that all tensors to save to.
If not specified, will use the model name.
size_threshold: Threshold for size of data. Only when tensor's data is >= the size_threshold it will be converted
to external data. To convert every tensor with raw data to external data set size_threshold=0.
convert_attribute: If true, convert all tensors to external data
If false, convert only non-attribute tensors to external data
'''
if isinstance(proto, bytes):
proto = _deserialize(proto, ModelProto())
if save_as_external_data:
convert_model_to_external_data(proto, all_tensors_to_one_file,
location, size_threshold,
convert_attribute)
model_filepath = _get_file_path(f)
if model_filepath:
basepath = os.path.dirname(model_filepath)
proto = write_external_data_tensors(proto, basepath)
s = _serialize(proto)
_save_bytes(s, f)
def export_model(sym, params, in_shapes=None, in_types=np.float32,
onnx_file_path='model.onnx', verbose=False, dynamic=False,
dynamic_input_shapes=None, run_shape_inference=False, input_type=None,
input_shape=None, large_model=False):
"""Exports the MXNet model file, passed as a parameter, into ONNX model.
Accepts both symbol,parameter objects as well as json and params filepaths as input.
Operator support and coverage -
https://github.com/apache/incubator-mxnet/tree/v1.x/python/mxnet/onnx#operator-support-matrix
Parameters
----------
sym : str or symbol object
Path to the json file or Symbol object
params : str or dict or list of dict
str - Path to the params file
dict - params dictionary (Including both arg_params and aux_params)
list - list of length 2 that contains arg_params and aux_params
in_shapes : List of tuple
Input shape of the model e.g [(1,3,224,224)]
in_types : data type or list of data types
Input data type e.g. np.float32, or [np.float32, np.int32]
onnx_file_path : str
Path where to save the generated onnx file
verbose : Boolean
If True will print logs of the model conversion
dynamic: Boolean
If True will allow for dynamic input shapes to the model
dynamic_input_shapes: list of tuple
Specifies the dynamic input_shapes. If None then all dimensions are set to None
run_shape_inference : Boolean
If True will run shape inference on the model
input_type : data type or list of data types
This is the old name of in_types. We keep this parameter name for backward compatibility
input_shape : List of tuple
This is the old name of in_shapes. We keep this parameter name for backward compatibility
large_model : Boolean
Whether to export a model that is larger than 2 GB. If true will save param tensors in separate
files along with .onnx model file. This feature is supported since onnx 1.8.0
Returns
-------
onnx_file_path : str
Onnx file path
Notes
-----
This method is available when you ``import mxnet.onnx``
"""
try:
import onnx
from onnx import helper, mapping, shape_inference
from onnx.defs import onnx_opset_version
except ImportError:
raise ImportError("Onnx and protobuf need to be installed. "
+ "Instructions to install - https://github.com/onnx/onnx")
if input_type is not None:
in_types = input_type
if input_shape is not None:
in_shapes = input_shape
converter = MXNetGraph()
opset_version = onnx_opset_version()
if not isinstance(in_types, list):
in_types = [in_types for _ in range(len(in_shapes))]
in_types_t = [mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(i_t)] for i_t in in_types]
assert len(in_types) == len(in_shapes), "The lengths of in_types and in_shapes must equal"
# if input parameters are strings(file paths), load files and create symbol parameter objects
if isinstance(sym, string_types) and isinstance(params, string_types):
logging.info("Converting json and weight file to sym and params")
sym_obj, params_obj = load_module(sym, params)
onnx_graph = converter.create_onnx_graph_proto(sym_obj, params_obj, in_shapes,
in_types_t,
verbose=verbose, opset_version=opset_version,
dynamic=dynamic, dynamic_input_shapes=dynamic_input_shapes)
elif isinstance(sym, symbol.Symbol) and isinstance(params, dict):
onnx_graph = converter.create_onnx_graph_proto(sym, params, in_shapes,
in_types_t,
verbose=verbose, opset_version=opset_version,
dynamic=dynamic, dynamic_input_shapes=dynamic_input_shapes)
elif isinstance(sym, symbol.Symbol) and isinstance(params, list) and len(params) == 2:
# when params contains arg_params and aux_params
p = {}
p.update(params[0])
p.update(params[1])
onnx_graph = converter.create_onnx_graph_proto(sym, p, in_shapes,
in_types_t,
verbose=verbose, opset_version=opset_version,
dynamic=dynamic, dynamic_input_shapes=dynamic_input_shapes)
else:
raise ValueError("Input sym and params should either be files or objects")
# Create the model (ModelProto)
onnx_model = helper.make_model(onnx_graph)
# Run shape inference on the model. Due to ONNX bug/incompatibility this may or may not crash
if run_shape_inference:
try:
onnx_model = shape_inference.infer_shapes(onnx_model)
except: # pylint: disable=bare-except
logging.info("Shape inference failed, original export is kept.")
if large_model:
from onnx.external_data_helper import convert_model_to_external_data
convert_model_to_external_data(onnx_model, all_tensors_to_one_file=False, location=onnx_file_path+'.data')
onnx.save_model(onnx_model, onnx_file_path)
onnx.checker.check_model(onnx_file_path)
return onnx_file_path
