path: root/compiler/one-cmds/one-import-onnx

#!/usr/bin/env bash
''''export SCRIPT_PATH="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)" # '''
''''export PY_PATH=${SCRIPT_PATH}/venv/bin/python                                       # '''
''''test -f ${PY_PATH} && exec ${PY_PATH} "$0" "$@"                                     # '''
''''echo "Error: Virtual environment not found. Please run 'one-prepare-venv' command." # '''
''''exit 255                                                                            # '''

# Copyright (c) 2021 Samsung Electronics Co., Ltd. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import sys
import tempfile
import onnx
import onnx_tf

# ONNX legalizer is an optional feature
# It enables conversion of some operations, but in experimental phase for now
try:
    import onnx_legalizer
    _onnx_legalizer_enabled = True
except ImportError:
    _onnx_legalizer_enabled = False

import onelib.make_cmd as _make_cmd
import onelib.utils as oneutils

# TODO Find better way to suppress trackback on error
sys.tracebacklimit = 0


def get_driver_cfg_section():
    return "one-import-onnx"


def _get_parser():
    parser = argparse.ArgumentParser(
        description='command line tool to convert ONNX to circle')

    oneutils.add_default_arg(parser)

    ## tf2tfliteV2 arguments
    tf2tfliteV2_group = parser.add_argument_group('converter arguments')

    # input and output path.
    tf2tfliteV2_group.add_argument(
        '-i', '--input_path', type=str, help='full filepath of the input file')
    tf2tfliteV2_group.add_argument(
        '-o', '--output_path', type=str, help='full filepath of the output file')

    # input and output arrays.
    tf2tfliteV2_group.add_argument(
        '-I',
        '--input_arrays',
        type=str,
        help='names of the input arrays, comma-separated')
    tf2tfliteV2_group.add_argument(
        '-O',
        '--output_arrays',
        type=str,
        help='names of the output arrays, comma-separated')

    # fixed options
    tf2tfliteV2_group.add_argument('--model_format', default='saved_model')
    tf2tfliteV2_group.add_argument('--converter_version', default='v2')

    parser.add_argument('--unroll_rnn', action='store_true', help='Unroll RNN operators')
    parser.add_argument(
        '--unroll_lstm', action='store_true', help='Unroll LSTM operators')
    parser.add_argument(
        '--keep_io_order',
        action='store_true',
        help=
        'Ensure generated circle model preserves the I/O order of the original onnx model.'
    )

    # save intermediate file(s)
    parser.add_argument(
        '--save_intermediate',
        action='store_true',
        help='Save intermediate files to output folder')

    # experimental options
    parser.add_argument(
        '--experimental_disable_batchmatmul_unfold',
        action='store_true',
        help='Experimental disable BatchMatMul unfold')

    return parser


def _verify_arg(parser, args):
    """verify given arguments"""
    # check if required arguments is given
    missing = []
    if not oneutils.is_valid_attr(args, 'input_path'):
        missing.append('-i/--input_path')
    if not oneutils.is_valid_attr(args, 'output_path'):
        missing.append('-o/--output_path')
    if len(missing):
        parser.error('the following arguments are required: ' + ' '.join(missing))


def _parse_arg(parser):
    args = parser.parse_args()
    # print version
    if args.version:
        oneutils.print_version_and_exit(__file__)

    return args


def _apply_verbosity(verbosity):
    # NOTE
    # TF_CPP_MIN_LOG_LEVEL
    #   0 : INFO + WARNING + ERROR + FATAL
    #   1 : WARNING + ERROR + FATAL
    #   2 : ERROR + FATAL
    #   3 : FATAL
    if verbosity:
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
    else:
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


# The index of input/output is added in front of the name. For example,
# Original input names: 'a', 'c', 'b'
# Renamed: '0001_a', '0002_c', '0003_b'
# This will preserve I/O order after import.
def _remap_io_names(onnx_model):
    # gather existing name of I/O and generate new name of I/O in sort order
    input_nodes = []
    output_nodes = []
    remap_inputs = []
    remap_outputs = []
    initializers = []
    # some models may have initializers as inputs. ignore them.
    for initializer in onnx_model.graph.initializer:
        initializers.append(initializer.name)
    for idx in range(0, len(onnx_model.graph.input)):
        name = onnx_model.graph.input[idx].name
        if not name in initializers:
            input_nodes.append(name)
            remap_inputs.append(format(idx + 1, '04d') + '_' + name)
    for idx in range(0, len(onnx_model.graph.output)):
        name = onnx_model.graph.output[idx].name
        output_nodes.append(name)
        remap_outputs.append(format(idx + 1, '04d') + '_' + name)
    # change names for graph input
    for i in range(len(onnx_model.graph.input)):
        if onnx_model.graph.input[i].name in input_nodes:
            to_rename = onnx_model.graph.input[i].name
            idx = input_nodes.index(to_rename)
            onnx_model.graph.input[i].name = remap_inputs[idx]
    # change names of all nodes in the graph
    for i in range(len(onnx_model.graph.node)):
        # check node.input is to change to remap_inputs or remap_outputs
        for j in range(len(onnx_model.graph.node[i].input)):
            if onnx_model.graph.node[i].input[j] in input_nodes:
                to_rename = onnx_model.graph.node[i].input[j]
                idx = input_nodes.index(to_rename)
                onnx_model.graph.node[i].input[j] = remap_inputs[idx]
            if onnx_model.graph.node[i].input[j] in output_nodes:
                to_rename = onnx_model.graph.node[i].input[j]
                idx = output_nodes.index(to_rename)
                onnx_model.graph.node[i].input[j] = remap_outputs[idx]
        # check node.output is to change to remap_inputs or remap_outputs
        for j in range(len(onnx_model.graph.node[i].output)):
            if onnx_model.graph.node[i].output[j] in output_nodes:
                to_rename = onnx_model.graph.node[i].output[j]
                idx = output_nodes.index(to_rename)
                onnx_model.graph.node[i].output[j] = remap_outputs[idx]
            if onnx_model.graph.node[i].output[j] in input_nodes:
                to_rename = onnx_model.graph.node[i].output[j]
                idx = input_nodes.index(to_rename)
                onnx_model.graph.node[i].output[j] = remap_inputs[idx]
    # change names for graph output
    for i in range(len(onnx_model.graph.output)):
        if onnx_model.graph.output[i].name in output_nodes:
            to_rename = onnx_model.graph.output[i].name
            idx = output_nodes.index(to_rename)
            onnx_model.graph.output[i].name = remap_outputs[idx]


def _convert(args):
    _apply_verbosity(args.verbose)

    # get file path to log
    dir_path = os.path.dirname(os.path.realpath(__file__))
    logfile_path = os.path.realpath(args.output_path) + '.log'

    with open(logfile_path, 'wb') as f, tempfile.TemporaryDirectory() as tmpdir:
        # save intermediate
        if oneutils.is_valid_attr(args, 'save_intermediate'):
            tmpdir = os.path.dirname(logfile_path)
        # convert onnx to tf saved model
        onnx_model = onnx.load(getattr(args, 'input_path'))
        if _onnx_legalizer_enabled:
            options = onnx_legalizer.LegalizeOptions
            options.unroll_rnn = oneutils.is_valid_attr(args, 'unroll_rnn')
            options.unroll_lstm = oneutils.is_valid_attr(args, 'unroll_lstm')
            onnx_legalizer.legalize(onnx_model, options)
        if oneutils.is_valid_attr(args, 'keep_io_order'):
            _remap_io_names(onnx_model)
            if oneutils.is_valid_attr(args, 'save_intermediate'):
                basename = os.path.basename(getattr(args, 'input_path'))
                fixed_path = os.path.join(tmpdir,
                                          os.path.splitext(basename)[0] + '~.onnx')
                onnx.save(onnx_model, fixed_path)
        tf_savedmodel = onnx_tf.backend.prepare(onnx_model)

        savedmodel_name = os.path.splitext(os.path.basename(
            args.output_path))[0] + '.savedmodel'
        savedmodel_output_path = os.path.join(tmpdir, savedmodel_name)
        tf_savedmodel.export_graph(savedmodel_output_path)

        # make a command to convert from tf to tflite
        tf2tfliteV2_path = os.path.join(dir_path, 'tf2tfliteV2.py')
        tf2tfliteV2_output_name = os.path.splitext(os.path.basename(
            args.output_path))[0] + '.tflite'
        tf2tfliteV2_output_path = os.path.join(tmpdir, tf2tfliteV2_output_name)

        tf2tfliteV2_cmd = _make_cmd.make_tf2tfliteV2_cmd(
            args, tf2tfliteV2_path, savedmodel_output_path, tf2tfliteV2_output_path)

        f.write((' '.join(tf2tfliteV2_cmd) + '\n').encode())

        # convert tf to tflite
        oneutils.run(tf2tfliteV2_cmd, logfile=f)

        # make a command to convert from tflite to circle
        tflite2circle_path = os.path.join(dir_path, 'tflite2circle')
        tflite2circle_cmd = _make_cmd.make_tflite2circle_cmd(tflite2circle_path,
                                                             tf2tfliteV2_output_path,
                                                             getattr(args, 'output_path'))

        f.write((' '.join(tflite2circle_cmd) + '\n').encode())

        # convert tflite to circle
        oneutils.run(tflite2circle_cmd, err_prefix="tflite2circle", logfile=f)


def main():
    # parse arguments
    parser = _get_parser()
    args = _parse_arg(parser)

    # parse configuration file
    oneutils.parse_cfg(args.config, 'one-import-onnx', args)

    # verify arguments
    _verify_arg(parser, args)

    # convert
    _convert(args)


if __name__ == '__main__':
    oneutils.safemain(main, __file__)