|
| 1 | +#!/usr/bin/env python |
| 2 | +# -*- coding: utf-8 -*- |
| 3 | +# |
| 4 | +# Copyright (c) 2022 Intel Corporation |
| 5 | +# |
| 6 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | +# you may not use this file except in compliance with the License. |
| 8 | +# You may obtain a copy of the License at |
| 9 | +# |
| 10 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | +# |
| 12 | +# Unless required by applicable law or agreed to in writing, software |
| 13 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | +# See the License for the specific language governing permissions and |
| 16 | +# limitations under the License. |
| 17 | +"""Helper functions to export model from TensorFlow to ONNX.""" |
| 18 | + |
| 19 | +import re |
| 20 | + |
| 21 | +from neural_compressor.utils import logger |
| 22 | +from neural_compressor.utils.utility import LazyImport |
| 23 | + |
| 24 | +t2o = LazyImport("tf2onnx") |
| 25 | + |
| 26 | + |
| 27 | +def _split_nodename_and_shape(name): |
| 28 | + """Split input name with shape into name and shape.""" |
| 29 | + # pattern for a node name |
| 30 | + inputs = [] |
| 31 | + shapes = {} |
| 32 | + # input takes in most cases the format name:0, where 0 is the output number |
| 33 | + # in some cases placeholders don't have a rank which onnx can't handle so we let uses override the shape |
| 34 | + # by appending the same, ie : [1,28,28,3] |
| 35 | + name_pattern = r"(?:([\w\d/\-\._:]+)(\[[\-\d,]+\])?),?" |
| 36 | + splits = re.split(name_pattern, name) |
| 37 | + for i in range(1, len(splits), 3): |
| 38 | + inputs.append(splits[i] + ":0") |
| 39 | + if splits[i + 1] is not None: |
| 40 | + shape = [int(n) for n in splits[i + 1][1:-1].split(",")] |
| 41 | + shape = [n if n >= 0 else None for n in shape] |
| 42 | + shapes[splits[i] + ":0"] = shape |
| 43 | + if not shapes: |
| 44 | + shapes = None |
| 45 | + return inputs, shapes |
| 46 | + |
| 47 | + |
| 48 | +def tf_to_fp32_onnx(graph_def, save_path, opset_version=14, input_names=None, output_names=None, inputs_as_nchw=None): |
| 49 | + """Export FP32 Tensorflow model into FP32 ONNX model using tf2onnx tool. |
| 50 | +
|
| 51 | + Args: |
| 52 | + graph_def (graph_def to convert): fp32 graph_def. |
| 53 | + save_path (str): save path of ONNX model. |
| 54 | + opset_version (int, optional): opset version. Defaults to 14. |
| 55 | + input_names (list, optional): input names. Defaults to None. |
| 56 | + output_names (list, optional): output names. Defaults to None. |
| 57 | + inputs_as_nchw (list, optional): transpose the input. Defaults to None. |
| 58 | + """ |
| 59 | + shape_override = None |
| 60 | + if isinstance(input_names, str): |
| 61 | + input_names, shape_override = _split_nodename_and_shape(input_names) |
| 62 | + else: |
| 63 | + input_names[:] = [o + ":0" for o in input_names] |
| 64 | + output_names[:] = [o + ":0" for o in output_names] |
| 65 | + t2o.convert.from_graph_def( |
| 66 | + graph_def=graph_def, |
| 67 | + input_names=input_names, |
| 68 | + output_names=output_names, |
| 69 | + inputs_as_nchw=inputs_as_nchw, |
| 70 | + shape_override=shape_override, |
| 71 | + opset=opset_version, |
| 72 | + output_path=save_path, |
| 73 | + ) |
| 74 | + info = "The FP32 ONNX Model exported to path: {0}".format(save_path) |
| 75 | + logger.info("*" * len(info)) |
| 76 | + logger.info(info) |
| 77 | + logger.info("*" * len(info)) |
| 78 | + |
| 79 | + |
| 80 | +def tf_to_int8_onnx( |
| 81 | + int8_model, save_path, opset_version: int = 14, input_names=None, output_names=None, inputs_as_nchw=None |
| 82 | +): |
| 83 | + """Export INT8 Tensorflow model into INT8 ONNX model. |
| 84 | +
|
| 85 | + Args: |
| 86 | + int8_model (tensorflow ITEX QDQ model): int8 model. |
| 87 | + save_path (str): save path of ONNX model. |
| 88 | + opset_version (int, optional): opset version. Defaults to 14. |
| 89 | + input_names (list, optional): input names. Defaults to None. |
| 90 | + output_names (list, optional): output names. Defaults to None. |
| 91 | + inputs_as_nchw (list, optional): transpose the input. Defaults to None. |
| 92 | + """ |
| 93 | + shape_override = None |
| 94 | + if isinstance(input_names, str): |
| 95 | + input_names, shape_override = _split_nodename_and_shape(input_names) |
| 96 | + else: |
| 97 | + input_names[:] = [o + ":0" for o in input_names] |
| 98 | + output_names[:] = [o + ":0" for o in output_names] |
| 99 | + onnx_convert_graph = "./converted_graph.onnx" |
| 100 | + from neural_compressor.adaptor.tf_utils.tf2onnx_converter import TensorflowQDQToOnnxQDQConverter |
| 101 | + |
| 102 | + TensorflowQDQToOnnxQDQConverter( |
| 103 | + int8_model, input_names, output_names, shape_override, inputs_as_nchw, opset_version |
| 104 | + ).convert(onnx_convert_graph) |
| 105 | + |
| 106 | + import onnxruntime as ort |
| 107 | + |
| 108 | + sess_options = ort.SessionOptions() |
| 109 | + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| 110 | + sess_options.optimized_model_filepath = save_path |
| 111 | + import onnx |
| 112 | + |
| 113 | + model = onnx.load(onnx_convert_graph) |
| 114 | + ort.InferenceSession(model.SerializeToString(), sess_options) |
| 115 | + info = "The INT8 ONNX Model is exported to path: {0}".format(save_path) |
| 116 | + logger.info("*" * len(info)) |
| 117 | + logger.info(info) |
| 118 | + logger.info("*" * len(info)) |
0 commit comments