def test_save_and_load_model(self): # type: () -> None proto = self._simple_model() cls = ModelProto proto_string = onnx._serialize(proto) # Test if input is string loaded_proto = onnx.load_model_from_string(proto_string) self.assertTrue(proto == loaded_proto) # Test if input has a read function f = io.BytesIO() onnx.save_model(proto_string, f) f = io.BytesIO(f.getvalue()) loaded_proto = onnx.load_model(f, cls) self.assertTrue(proto == loaded_proto) # Test if input is a file name try: fi = tempfile.NamedTemporaryFile(delete=False) onnx.save_model(proto, fi) fi.close() loaded_proto = onnx.load_model(fi.name, cls) self.assertTrue(proto == loaded_proto) finally: os.remove(fi.name)
def save_model_to_file(self, output_path, use_external_data_format=False): ''' Save model to external data, which is needed for model size > 2GB ''' if use_external_data_format: onnx.external_data_helper.convert_model_to_external_data(self.model, all_tensors_to_one_file=True, location=Path(output_path).name + ".data") onnx.save_model(self.model, output_path)
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 onnx_convert(keras_model_file, output_file): custom_object_dict = get_custom_objects() model = load_model(keras_model_file, custom_objects=custom_object_dict) # convert to onnx model onnx_model = keras2onnx.convert_keras(model, model.name) # save converted onnx model onnx.save_model(onnx_model, output_file)
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 quantize_dynamic(model_input: Path, model_output: Path, op_types_to_quantize=[], per_channel=False, reduce_range=False, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], nodes_to_exclude=[]): ''' Given an onnx model, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize :param model_output: file path of quantized model :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default :param per_channel: quantize weights per channel :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode :param nbits: number of bits to represent quantized data. Currently only supporting 8-bit types :param activation_type: quantization data type of activation :param weight_type: quantization data type of weight :param nodes_to_quantize: List of nodes names to quantize. When this list is not None only the nodes in this list are quantized. example: [ 'Conv__224', 'Conv__252' ] :param nodes_to_exclude: List of nodes names to exclude. The nodes in this list will be excluded from quantization when it is not None. ''' input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 mode = QuantizationMode.IntegerOps #optimize the original model optimized_model = optimize_model(Path(model_input)) if not op_types_to_quantize or len(op_types_to_quantize) == 0: op_types_to_quantize = list(IntegerOpsRegistry.keys()) quantizer = ONNXQuantizer( optimized_model, per_channel, reduce_range, mode, False, #static weight_qType, input_qType, None, nodes_to_quantize, nodes_to_exclude, op_types_to_quantize) quantizer.quantize_model() onnx.save_model(quantizer.model.model, model_output)
def onnx_convert(keras_model_file, output_file, op_set): #custom_object_dict = get_custom_objects() model = load_model(keras_model_file) # convert to onnx model onnx_model = keras2onnx.convert_keras(model, model.name, target_opset=op_set) # save converted onnx model onnx.save_model(onnx_model, output_file)
def cpu_infer(self, case_dir: str, model_file: bytes, type: str, mode: str): # create session try: print('[onnx]: using simplified model') sess = ort.InferenceSession(model_file) except Exception as e: print(e) try: print('[onnx]: using origin model') model_file = os.path.join(case_dir, 'test.onnx') sess = ort.InferenceSession(model_file) except Exception as e: print(e) print('[onnx]: using converted model') onnx_model = onnx.load(model_file) onnx_model = version_converter.convert_version(onnx_model, 8) model_file = os.path.join(case_dir, 'converted.onnx') onnx.save_model(onnx_model, model_file) sess = ort.InferenceSession(model_file) if mode is "dataset": for input in self.inputs: input_dict = {} for in_data in input['data']: topk = [] input_dict[input['name']] = self.transform_input( self.data_pre_process(in_data[0]), "float32", "CPU") outputs = sess.run(None, input_dict) for output in outputs: topk.append((in_data[1], get_topK('cpu', 1, output)[0])) if os.path.exists(os.path.join(case_dir, "cpu_dataset.txt")): os.remove(os.path.join(case_dir, "cpu_dataset.txt")) self.output_paths.append( (os.path.join(case_dir, 'cpu_result_0.bin'), os.path.join(case_dir, 'cpu_result_0.txt'))) with open(self.output_paths[-1][1], 'a') as f: for i in range(len(topk)): f.write(topk[i][0].split('/')[-1] + " " + str(topk[i][1]) + '\n') else: input_dict = {} for input in self.inputs: input_dict[input['name']] = self.transform_input( self.data_pre_process(input['data']), "float32", "CPU") outputs = sess.run(None, input_dict) i = 0 for output in outputs: bin_file = os.path.join(case_dir, f'cpu_result_{i}.bin') text_file = os.path.join(case_dir, f'cpu_result_{i}.txt') self.output_paths.append((bin_file, text_file)) output.tofile(bin_file) self.totxtfile(text_file, output) i += 1
def save_model_to_file(self, output_path): logger.info(f"Output model to {output_path}") if output_path.endswith(".json"): assert isinstance(self.model, ModelProto) with open(output_path, "w") as out: out.write(str(self.model)) else: save_model(self.model, output_path, format=None)
def setup_module(): onnx.save_model(create_onnx_model(), onnx_model_filename) onnx.save_model(create_onnx_model_2(), onnx_model_2_filename) onnx.save_model(create_onnx_model_with_custom_attributes(), onnx_model_with_custom_attributes_filename) onnx.save_model(create_onnx_model_with_subgraphs(), onnx_model_with_subgraphs_filename) onnx.save_model(create_onnx_model_for_op_extension(), onnx_model_for_op_extension_test)
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 quantize(onnx_model_path: Path) -> Path: """ Quantize the weights of the model from float32 to in8 to allow very efficient inference on modern CPU Args: onnx_model_path: Path to location the exported ONNX model is stored Returns: The Path generated for the quantized """ import onnx from onnx.onnx_pb import ModelProto from onnxruntime.quantization import QuantizationMode from onnxruntime.quantization.onnx_quantizer import ONNXQuantizer from onnxruntime.quantization.registry import IntegerOpsRegistry # Load the ONNX model onnx_model = onnx.load(onnx_model_path.as_posix()) if parse(onnx.__version__) < parse("1.5.0"): print( "Models larger than 2GB will fail to quantize due to protobuf constraint.\n" "Please upgrade to onnxruntime >= 1.5.0.") # Copy it copy_model = ModelProto() copy_model.CopyFrom(onnx_model) # Construct quantizer quantizer = ONNXQuantizer( model=copy_model, per_channel=False, reduce_range=False, mode=QuantizationMode.IntegerOps, static=False, weight_qType=True, input_qType=False, tensors_range=None, nodes_to_quantize=None, nodes_to_exclude=None, op_types_to_quantize=list(IntegerOpsRegistry), ) # Quantize and export quantizer.quantize_model() # Append "-quantized" at the end of the model's name quantized_model_path = generate_identified_filename( onnx_model_path, "-quantized") # Save model print( f"Quantized model has been written at {quantized_model_path}: \N{heavy check mark}" ) onnx.save_model(quantizer.model.model, quantized_model_path.as_posix()) return quantized_model_path
def save_model(onnx_model, path, conda_env=None, mlflow_model=Model()): """ Save an ONNX model to a path on the local file system. :param onnx_model: ONNX model to be saved. :param path: Local path where the model is to be saved. :param conda_env: Either a dictionary representation of a Conda environment or the path to a Conda environment yaml file. If provided, this describes the environment this model should be run in. At minimum, it should specify the dependencies contained in :func:`get_default_conda_env()`. If `None`, the default :func:`get_default_conda_env()` environment is added to the model. The following is an *example* dictionary representation of a Conda environment:: { 'name': 'mlflow-env', 'channels': ['defaults'], 'dependencies': [ 'python=3.6.0', 'onnx=1.4.1', 'onnxruntime=0.3.0' ] } :param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to. """ import onnx path = os.path.abspath(path) if os.path.exists(path): raise MlflowException(message="Path '{}' already exists".format(path), error_code=RESOURCE_ALREADY_EXISTS) os.makedirs(path) model_data_subpath = "model.onnx" model_data_path = os.path.join(path, model_data_subpath) # Save onnx-model onnx.save_model(onnx_model, model_data_path) conda_env_subpath = "conda.yaml" if conda_env is None: conda_env = get_default_conda_env() elif not isinstance(conda_env, dict): with open(conda_env, "r") as f: conda_env = yaml.safe_load(f) with open(os.path.join(path, conda_env_subpath), "w") as f: yaml.safe_dump(conda_env, stream=f, default_flow_style=False) pyfunc.add_to_model(mlflow_model, loader_module="mlflow.onnx", data=model_data_subpath, env=conda_env_subpath) mlflow_model.add_flavor(FLAVOR_NAME, onnx_version=onnx.__version__, data=model_data_subpath) mlflow_model.save(os.path.join(path, "MLmodel"))
def test_compile_with_no_input_formats(self): with NamedTemporaryFile(suffix='.onnx') as model_file: onnx.save_model(_make_onnx_model().model_proto, model_file.name) config = Config.from_json({}) compiled = onnx_compiler.compile_source(source=ONNXModelFile( model_file.name), config=config) self.assertEqual(compiled.input_data_formats, [None, None])
def export_and_recurse(node, attribute, output_dir, level): name = node.name name = name.replace('/', '_') sub_model = onnx.ModelProto() sub_model.graph.MergeFrom(attribute.g) filename = 'L' + str( level ) + '_' + node.op_type + '_' + attribute.name + '_' + name + '.onnx' onnx.save_model(sub_model, os.path.join(output_dir, filename)) dump_subgraph(sub_model, output_dir, level + 1)
def export_and_recurse(node, attribute, output_dir, level): name = node.name name = name.replace("/", "_") sub_model = onnx.ModelProto() sub_model.graph.MergeFrom(attribute.g) filename = "L" + str( level ) + "_" + node.op_type + "_" + attribute.name + "_" + name + ".onnx" onnx.save_model(sub_model, os.path.join(output_dir, filename)) dump_subgraph(sub_model, output_dir, level + 1)
def run_onnx_runtime(self, case_name, onnx_model, data, expected, rtol=1.e-3, atol=1.e-6): temp_model_file = TestKerasApplications.get_temp_file('temp_' + case_name + '.onnx') onnx.save_model(onnx_model, temp_model_file) try: import onnxruntime sess = onnxruntime.InferenceSession(temp_model_file) except ImportError: return True if not isinstance(expected, list): expected = [expected] data = data if isinstance(data, list) else [data] input_names = sess.get_inputs() # to avoid too complicated test code, we restrict the input name in Keras test cases must be # in alphabetical order. It's always true unless there is any trick preventing that. feed = zip(sorted(i_.name for i_ in input_names), data) actual = sess.run(None, dict(feed)) res = all( np.allclose(expected[n_], actual[n_], rtol=rtol, atol=atol) for n_ in range(len(expected))) if res and temp_model_file not in self.model_files: # still keep the failed case files for the diagnosis. self.model_files.append(temp_model_file) if not res: for n_ in range(len(expected)): expected_list = expected[n_].flatten() actual_list = actual[n_].flatten() diff_list = abs(expected_list - actual_list) count_total = len(expected_list) count_error = 0 for e_, a_, d_ in zip(expected_list, actual_list, diff_list): if d_ > atol + rtol * abs(a_): if count_error < 10: # print the first 10 mismatches print("case = " + case_name + ", result mismatch for expected = " + str(e_) + ", actual = " + str(a_), file=sys.stderr) count_error = count_error + 1 print("case = " + case_name + ", " + str(count_error) + " mismatches out of " + str(count_total) + " for list " + str(n_), file=sys.stderr) return res
def SaveOnnxModel(onnx_model, onnx_save_path, need_polish=True): try: if need_polish: polished_model = onnx.utils.polish_model(onnx_model) onnx.save_model(polished_model, onnx_save_path) else: onnx.save_model(onnx_model, onnx_save_path) print("模型保存成功,已保存至:" + onnx_save_path) except Exception as e: print("模型存在问题,未保存成功:", e)
def test_mask_rcnn(self): onnx_model = keras2onnx.convert_keras( model.keras_model, target_opset=10, custom_op_conversions=tf2onnx_contrib_op_conversion) import skimage img_path = os.path.join(os.path.dirname(__file__), '../data', 'street.jpg') image = skimage.io.imread(img_path) images = [image] case_name = 'mask_rcnn' if not os.path.exists(tmp_path): os.mkdir(tmp_path) temp_model_file = os.path.join(tmp_path, 'temp_' + case_name + '.onnx') onnx.save_model(onnx_model, temp_model_file) try: import onnxruntime sess = onnxruntime.InferenceSession(temp_model_file) except ImportError: return True # preprocessing molded_images, image_metas, windows = model.mold_inputs(images) anchors = model.get_anchors(molded_images[0].shape) anchors = np.broadcast_to(anchors, (model.config.BATCH_SIZE, ) + anchors.shape) expected = model.keras_model.predict([ molded_images.astype(np.float32), image_metas.astype(np.float32), anchors ]) actual = \ sess.run(None, {"input_image": molded_images.astype(np.float32), "input_anchors": anchors, "input_image_meta": image_metas.astype(np.float32)}) rtol = 1.e-3 atol = 1.e-6 compare_idx = [0, 3] res = all( np.allclose(expected[n_], actual[n_], rtol=rtol, atol=atol) for n_ in compare_idx) if res and temp_model_file not in self.model_files: # still keep the failed case files for the diagnosis. self.model_files.append(temp_model_file) if not res: for n_ in compare_idx: expected_list = expected[n_].flatten() actual_list = actual[n_].flatten() print_mismatches(case_name, n_, expected_list, actual_list, atol, rtol) self.assertTrue(res)
def test_to_array_with_external_data(self): # type: () -> None onnx.save_model(self.model, self.model_file_path, save_as_external_data=True, all_tensors_to_one_file=False, size_threshold=0) # raw_data of external tensor is not loaded model = onnx.load(self.model_file_path, load_external_data=False) # Specify self.temp_dir to load external tensor loaded_large_data = to_array(model.graph.initializer[0], self.temp_dir) self.assertTrue(np.allclose(loaded_large_data, self.large_data))
def test_node_name_type_custom_functions(self): # type: () -> None def convert_acos(builder, node, graph, err): params = NeuralNetwork_pb2.CustomLayerParams() params.className = node.op_type params.description = "Custom layer that corresponds to the ONNX op {}".format(node.op_type, ) builder.add_custom( name=node.name, input_names=node.inputs, output_names=node.outputs, custom_proto_spec=params ) def convert_topk_generic(builder, node, graph, err): params = NeuralNetwork_pb2.CustomLayerParams() params.className = node.op_type params.description = "Custom layer that corresponds to the ONNX op {}".format(node.op_type, ) params.parameters["axis"].intValue = node.attrs.get('axis', -1) params.parameters["k"].intValue = node.attrs['k'] builder.add_custom( name=node.name, input_names=node.inputs, output_names=node.outputs, custom_proto_spec=params ) def convert_topk_node_specific(builder, node, graph, err): params = NeuralNetwork_pb2.CustomLayerParams() params.className = node.op_type params.description = "Custom layer that corresponds to the ONNX op {}".format(node.op_type, ) params.parameters["axis"].intValue = node.attrs.get('axis', -1) builder.add_custom( name=node.name, input_names=node.inputs, output_names=node.outputs, custom_proto_spec=params ) onnx_model = _make_model_acos_exp_topk() onnx.save_model(onnx_model, 'acos.onnx') coreml_model = convert(model=onnx_model, add_custom_layers=True, custom_conversion_functions={'Acos':convert_acos, 'TopK':convert_topk_generic, 'output_values_output_indices':convert_topk_node_specific}) spec = coreml_model.get_spec() layers = spec.neuralNetwork.layers self.assertIsNotNone(layers[0].custom) self.assertIsNotNone(layers[2].custom) self.assertEqual('Acos', layers[0].custom.className) self.assertEqual('TopK', layers[2].custom.className) self.assertEqual(0, layers[2].custom.parameters['axis'].intValue)
def test_save_model_does_not_convert_to_external_data_and_saves_the_model(self): # type: () -> None model_file_path = self.get_temp_model_filename() onnx.save_model(self.model, model_file_path, save_as_external_data=False) self.assertTrue(Path.isfile(model_file_path)) model = onnx.load_model(model_file_path) initializer_tensor = model.graph.initializer[0] self.assertFalse(initializer_tensor.HasField("data_location")) attribute_tensor = model.graph.node[0].attribute[0].t self.assertFalse(attribute_tensor.HasField("data_location"))
def main(): if len(sys.argv) < 4: print('decast.py model_in model_out <op1, ...>') return input = sys.argv[1] output = sys.argv[2] op_list = sys.argv[3:] oxml = onnx.load_model(input) oxml = decast(oxml, op_list) onnx.save_model(oxml, output)
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_weights(self, model, epoch): save_path = self.get_save_path(epoch) model.save_weights(save_path, overwrite=True) # try: if _has_onnx: save_path = self.get_save_path(epoch, ext='onnx') onnx_model = keras2onnx.convert_keras(model, model.name, target_opset=10) onnx.save_model(onnx_model, save_path) # except Exception as e: # print(e) return
def test_yolov3(self): img_path = os.path.join(os.path.dirname(__file__), '../data', 'street.jpg') yolo3_yolo3_dir = os.path.join(os.path.dirname(__file__), '../../../keras-yolo3/yolo3') try: import onnxruntime except ImportError: return True from PIL import Image for is_tiny_yolo in [True, False]: if is_tiny_yolo: if not os.path.exists(tiny_model_file_name): urllib.request.urlretrieve(YOLOV3_TINY_WEIGHTS_PATH, tiny_model_file_name) yolo_weights = load_model(tiny_model_file_name) model_path = tiny_model_file_name # model path or trained weights path anchors_path = 'model_data/tiny_yolo_anchors.txt' case_name = 'yolov3-tiny' else: if not os.path.exists(model_file_name): urllib.request.urlretrieve(YOLOV3_WEIGHTS_PATH, model_file_name) yolo_weights = load_model(model_file_name) model_path = model_file_name # model path or trained weights path anchors_path = 'model_data/yolo_anchors.txt' case_name = 'yolov3' my_yolo = YOLO(model_path, anchors_path, yolo3_yolo3_dir) my_yolo.load_model(yolo_weights) onnx_model = convert_model(my_yolo, is_tiny_yolo) if not os.path.exists(tmp_path): os.mkdir(tmp_path) temp_model_file = os.path.join(tmp_path, 'temp_' + case_name + '.onnx') onnx.save_model(onnx_model, temp_model_file) sess = onnxruntime.InferenceSession(temp_model_file) image = Image.open(img_path) image_data = my_yolo.prepare_keras_data(image) all_boxes_k, all_scores_k, indices_k = my_yolo.final_model.predict([image_data, np.array([image.size[1], image.size[0]], dtype='float32').reshape(1, 2)]) image_data_onnx = np.transpose(image_data, [0, 3, 1, 2]) feed_f = dict(zip(['input_1', 'image_shape'], (image_data_onnx, np.array([image.size[1], image.size[0]], dtype='float32').reshape(1, 2)))) all_boxes, all_scores, indices = sess.run(None, input_feed=feed_f) expected = self.post_compute(all_boxes_k, all_scores_k, indices_k) actual = self.post_compute(all_boxes, all_scores, indices) res = all(np.allclose(expected[n_], actual[n_]) for n_ in range(3)) self.assertTrue(res)
def test_yolov3(self): img_path = os.path.join(os.path.dirname(__file__), '../data', 'street.jpg') yolo3_yolo3_dir = os.path.join(os.path.dirname(__file__), '../../keras-yolo3/yolo3') if not os.path.exists(model_file_name): urllib.request.urlretrieve(YOLOV3_WEIGHTS_PATH, model_file_name) yolo_weights = load_model(model_file_name) my_yolo = YOLO(yolo3_yolo3_dir) my_yolo.load_model(yolo_weights) case_name = 'yolov3' target_opset = 10 onnx_model = keras2onnx.convert_keras(my_yolo.final_model, target_opset=target_opset, channel_first_inputs=['input_1']) if not os.path.exists(tmp_path): os.mkdir(tmp_path) temp_model_file = os.path.join(tmp_path, 'temp_' + case_name + '.onnx') onnx.save_model(onnx_model, temp_model_file) try: import onnxruntime sess = onnxruntime.InferenceSession(temp_model_file) except ImportError: return True from PIL import Image image = Image.open(img_path) image_data = my_yolo.prepare_keras_data(image) all_boxes_k, all_scores_k, indices_k = my_yolo.final_model.predict([ image_data, np.array([image.size[1], image.size[0]], dtype='float32').reshape(1, 2) ]) image_data_onnx = np.transpose(image_data, [0, 3, 1, 2]) feed_f = dict( zip(['input_1', 'image_shape'], (image_data_onnx, np.array([image.size[1], image.size[0]], dtype='float32').reshape(1, 2)))) all_boxes, all_scores, indices = sess.run(None, input_feed=feed_f) expected = self.post_compute(all_boxes_k, all_scores_k, indices_k) actual = self.post_compute(all_boxes, all_scores, indices) res = all(np.allclose(expected[n_], actual[n_]) for n_ in range(3)) self.assertTrue(res)
def main(args): if args.weight_h5 == '': print("Not found --weight_path!") exit() onnx_model_name = args.weight_onnx model = get_model(width=224, height=224, depth=32) model.load_weights(args.weight_h5) onnx_model = keras2onnx.convert_keras(model, model.name) onnx.save_model(onnx_model, onnx_model_name) print("Saved file:", args.weight_onnx)
def test_save_model_with_external_data_multiple_times( self): # type: () -> None # Test onnx.save should respectively handle typical tensor and external tensor properly # 1st save: save two tensors which have raw_data # Only w_large will be stored as external tensors since it's larger than 1024 onnx.save_model(self.model, self.model_file_path, save_as_external_data=True, all_tensors_to_one_file=False, location=None, size_threshold=1024, convert_attribute=True) model_without_loading_external = onnx.load(self.model_file_path, load_external_data=False) large_input_tensor = model_without_loading_external.graph.initializer[ 0] self.assertTrue(large_input_tensor.HasField("data_location")) self.assertTrue( np.allclose(to_array(large_input_tensor, self.temp_dir), self.large_data)) small_shape_tensor = model_without_loading_external.graph.initializer[ 1] self.assertTrue(not small_shape_tensor.HasField("data_location")) self.assertTrue( np.allclose(to_array(small_shape_tensor), self.small_data)) # 2nd save: one tensor has raw_data (small); one external tensor (large) # Save them both as external tensors this time onnx.save_model(model_without_loading_external, self.model_file_path, save_as_external_data=True, all_tensors_to_one_file=False, location=None, size_threshold=0, convert_attribute=True) model_without_loading_external = onnx.load(self.model_file_path, load_external_data=False) large_input_tensor = model_without_loading_external.graph.initializer[ 0] self.assertTrue(large_input_tensor.HasField("data_location")) self.assertTrue( np.allclose(to_array(large_input_tensor, self.temp_dir), self.large_data)) small_shape_tensor = model_without_loading_external.graph.initializer[ 1] self.assertTrue(small_shape_tensor.HasField("data_location")) self.assertTrue( np.allclose(to_array(small_shape_tensor, self.temp_dir), self.small_data))
def quantize_qat(model_input: Path, model_output: Path, op_types_to_quantize=[], per_channel=False, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], nodes_to_exclude=[]): ''' Given a quantize-aware traning onnx model, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize :param model_output: file path of quantized model :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default. :param per_channel: quantize weights per channel :param activation_type: quantization data type of activation :param nodes_to_quantize: List of nodes names to quantize. When this list is not None only the nodes in this list are quantized. example: [ 'Conv__224', 'Conv__252' ] :param nodes_to_exclude: List of nodes names to exclude. The nodes in this list will be excluded from quantization when it is not None. ''' input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 mode = QuantizationMode.IntegerOps #optimize the original model optimized_model = optimize_model(Path(model_input)) if not op_types_to_quantize or len(op_types_to_quantize) == 0: op_types_to_quantize = list(IntegerOpsRegistry.keys()) quantizer = ONNXQuantizer( optimized_model, per_channel, mode, False, #static weight_qType, input_qType, None, nodes_to_quantize, nodes_to_exclude, op_types_to_quantize) quantizer.quantize_model() onnx.save_model(quantizer.model.model, model_output)
def run_onnx_runtime(case_name, onnx_model, data, expected, model_files, rtol=1.e-3, atol=1.e-6): if not os.path.exists(tmp_path): os.mkdir(tmp_path) temp_model_file = os.path.join(tmp_path, 'temp_' + case_name + '.onnx') onnx.save_model(onnx_model, temp_model_file) try: import onnxruntime sess = onnxruntime.InferenceSession(temp_model_file) except ImportError: keras2onnx.common.k2o_logger().warning("Cannot import ONNXRuntime!") return True if isinstance(data, dict): feed_input = data else: data = data if isinstance(data, list) else [data] input_names = sess.get_inputs() # to avoid too complicated test code, we restrict the input name in Keras test cases must be # in alphabetical order. It's always true unless there is any trick preventing that. feed = zip(sorted(i_.name for i_ in input_names), data) feed_input = dict(feed) actual = sess.run(None, feed_input) if expected is None: return if isinstance(expected, tuple): expected = list(expected) elif not isinstance(expected, list): expected = [expected] res = all( np.allclose(expected[n_], actual[n_], rtol=rtol, atol=atol) for n_ in range(len(expected))) if res and temp_model_file not in model_files: # still keep the failed case files for the diagnosis. model_files.append(temp_model_file) if not res: for n_ in range(len(expected)): expected_list = expected[n_].flatten() actual_list = actual[n_].flatten() print_mismatches(case_name, n_, expected_list, actual_list, rtol, atol) return res