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