Exemplo n.º 1
0
        def get_custom_model_spec():
            from coremltools.models.neural_network import NeuralNetworkBuilder
            from coremltools.models.datatypes import Array, Dictionary, String

            input_name = 'output1'
            input_length = self._feature_extractor.output_length
            builder = NeuralNetworkBuilder(
                [(input_name, Array(input_length, ))],
                [(prob_name, Dictionary(String))], 'classifier')

            ctx = _mxnet_utils.get_mxnet_context()[0]
            input_name, output_name = input_name, 0
            import mxnet as _mx
            for i, cur_layer in enumerate(self._custom_classifier):
                output_name = str(i)
                if type(cur_layer) == _mx.gluon.nn.basic_layers.Dense:
                    W = cur_layer.weight.data(ctx).asnumpy()
                    nC, nB = W.shape
                    Wb = cur_layer.bias.data(ctx).asnumpy()

                    builder.add_inner_product(name='inner_product_' + str(i),
                                              W=W,
                                              b=Wb,
                                              input_channels=nB,
                                              output_channels=nC,
                                              has_bias=True,
                                              input_name=input_name,
                                              output_name='inner_product_' +
                                              output_name)
                    if cur_layer.act:
                        builder.add_activation("activation" + str(i), 'RELU',
                                               'inner_product_' + output_name,
                                               output_name)
                elif type(cur_layer) == _mx.gluon.nn.basic_layers.BatchNorm:
                    zeros = _np.zeros(nC)
                    ones = _np.ones(nC)
                    builder.add_batchnorm(name='bn_layer_' + str(i),
                                          channels=nC,
                                          gamma=ones,
                                          beta=zeros,
                                          mean=zeros,
                                          variance=ones,
                                          input_name=input_name,
                                          output_name=output_name)
                elif type(cur_layer) == _mx.gluon.nn.basic_layers.Dropout:
                    continue
                input_name = output_name

            last_output = builder.spec.neuralNetworkClassifier.layers[
                -1].output[0]
            builder.add_softmax('softmax', last_output, self.target)

            builder.set_class_labels(self.classes)
            builder.set_input([input_name], [(input_length, )])
            builder.set_output([self.target], [(self.num_classes, )])

            return builder.spec
        def get_custom_model_spec():
            from coremltools.models.neural_network import NeuralNetworkBuilder
            from coremltools.models.datatypes import Array

            input_name = "output1"
            input_length = self._feature_extractor.output_length
            builder = NeuralNetworkBuilder(
                [(input_name, Array(input_length,))],
                [(prob_name, Array(self.num_classes,))],
                "classifier",
            )
            layer_counter = [0]
            builder.set_input([input_name], [(input_length,)])

            def next_layer_name():
                layer_counter[0] += 1
                return "layer_%d" % layer_counter[0]

            for i, cur_layer in enumerate(self._custom_classifier.export_weights()):
                W = cur_layer["weight"]
                nC, nB = W.shape
                Wb = cur_layer["bias"]

                output_name = next_layer_name()
                builder.add_inner_product(
                    name="inner_product_" + str(i),
                    W=W,
                    b=Wb,
                    input_channels=nB,
                    output_channels=nC,
                    has_bias=True,
                    input_name=input_name,
                    output_name=output_name,
                )

                input_name = output_name

                if cur_layer["act"]:
                    output_name = next_layer_name()
                    builder.add_activation(
                        "activation" + str(i), "RELU", input_name, output_name
                    )
                    input_name = output_name

            builder.add_softmax("softmax", input_name, prob_name)
            builder.set_class_labels(
                self.classes,
                predicted_feature_name=self.target,
                prediction_blob=prob_name,
            )
            return builder.spec
Exemplo n.º 3
0
        def get_custom_model_spec():
            from coremltools.models.neural_network import NeuralNetworkBuilder
            from coremltools.models.datatypes import Array, Dictionary, String

            input_name = 'output1'
            input_length = self._feature_extractor.output_length
            builder = NeuralNetworkBuilder(
                [(input_name, Array(input_length, ))],
                [(prob_name, Dictionary(String))], 'classifier')

            input_name, output_name = input_name, 0
            for i, cur_layer in enumerate(
                    self._custom_classifier.export_weights()):
                W = cur_layer['weight']
                nC, nB = W.shape
                Wb = cur_layer['bias']

                builder.add_inner_product(name="inner_product_" + str(i),
                                          W=W,
                                          b=Wb,
                                          input_channels=nB,
                                          output_channels=nC,
                                          has_bias=True,
                                          input_name=str(input_name),
                                          output_name='inner_product_' +
                                          str(output_name))

                if cur_layer['act']:
                    builder.add_activation("activation" + str(i), 'RELU',
                                           'inner_product_' + str(output_name),
                                           str(output_name))

                input_name = i
                output_name = i + 1

            last_output = builder.spec.neuralNetworkClassifier.layers[
                -1].output[0]
            builder.add_softmax('softmax', last_output, self.target)

            builder.set_class_labels(self.classes,
                                     predicted_feature_name=self.target)
            builder.set_input([input_name], [(input_length, )])
            builder.set_output([self.target], [(self.num_classes, )])

            return builder.spec
Exemplo n.º 4
0
def _convert_pb_to_mlmodel(
        tf_model_path,
        mlmodel_path,
        output_feature_names,
        input_name_shape_dict={},
        image_input_names=None,
        is_bgr=False,
        red_bias=0.0,
        green_bias=0.0,
        blue_bias=0.0,
        gray_bias=0.0,
        image_scale=1.0,
        class_labels=None,
        predicted_feature_name=None,
        predicted_probabilities_output='',
        add_custom_layers=False,  # type: bool
        custom_conversion_functions={},  # type: Dict[Text, Any]
):

    # Load the TF graph
    print('')
    print('Loading the TF graph...')
    with open(tf_model_path, 'rb') as f:
        serialized = f.read()

    tf.reset_default_graph()
    gdef = tf.GraphDef()
    gdef.ParseFromString(serialized)

    with tf.Graph().as_default() as g:
        tf.import_graph_def(gdef, name='')

    sess = tf.Session(graph=g)
    OPS = g.get_operations()

    if 'DecodeJpeg' in [op.type for op in OPS]:
        raise NotImplementedError(
            "Unsupported Op of type: DecodeJpeg. "
            "Kindly refer to the \"examples/inception_v3.ipynb\" notebook, "
            "on the tfcoreml github page, to see how to strip input "
            "pre-processing from the TF graph before conversion to CoreML.")

    print('Graph Loaded.')
    # Sort the ops in topological order and check whether the graph has cycles, if yes, error out
    OPS = _topological_sort_ops(OPS)

    SHAPE_DICT = {}  #Tensor name --> shape ({str: list})
    CONSTS = {}  #Const Tensor name --> value
    BLOB_GRAPH = {}  #Blob name to list of ops it feeds into

    # Make Dictionary of Input blob to the list of ops it feeds into
    for op in OPS:
        for inp in op.inputs:
            if inp.name in BLOB_GRAPH:
                BLOB_GRAPH[inp.name].append(op)
        for out in op.outputs:
            if out.name not in BLOB_GRAPH:
                BLOB_GRAPH[out.name] = []

    # Fill in input information
    input_features = []
    output_features = []
    input_feed_dict = dict()  #Input tensors' values
    input_feed_dict2 = dict()  # used later to find skippable ops

    # run through all placeholders
    for op in OPS:
        output_names = set([compat.as_str_any(x.name) for x in op.outputs])
        if op.type == 'Placeholder':
            # Handle placeholders -- all placeholders are inputs
            assert not any(filter(output_names.__contains__, output_feature_names)), \
                ('Output feature cannot be a placeholder')
            input_name = compat.as_str_any(op.outputs[0].name)
            shape = op.outputs[0].get_shape()

            if input_name in input_name_shape_dict:
                shape = input_name_shape_dict[input_name]
            elif shape.is_fully_defined():
                shape = shape.as_list()
            else:
                try:
                    shape_list = shape.as_list()
                except:
                    raise ValueError(
                        'Please provide the shape for the input {} through the argument \'input_name_shape_dict\''
                        .format(input_name))
                if shape_list[0] is None and None not in shape_list[1:]:
                    shape = [1] + shape_list[1:]
                else:
                    raise ValueError(
                        "%s is a placeholder with incomplete shape %s. Please provide the 'input_name_shape_dict' "
                        "argument to the convert function, with the fully defined shape."
                        % (input_name, str(shape)))

            if len(shape) == 0:  # scalar - use a 1
                input_feed_dict[op.outputs[0]] = 1
                input_feed_dict2[op.outputs[0]] = 1
            else:
                input_feed_dict[op.outputs[0]] = np.random.rand(*shape)
                input_feed_dict2[op.outputs[0]] = 255 * np.random.rand(*shape)

            SHAPE_DICT[input_name] = list(shape)

    # Find "effectively_constant_ops": ops whose output(s) do not change with different valued Graph level inputs
    # Find "unused_ops" : ops that are not connected to the output(s)
    unused_ops = []
    effectively_constant_ops = []
    try:
        print(
            "Now finding ops in the TF graph that can be dropped for inference"
        )
        unused_ops, effectively_constant_ops = _find_unused_ops(
            OPS, sess, output_feature_names, input_feed_dict,
            input_feed_dict2)  # return type: List[str], List[str]
    except:
        pass

    # Populate SHAPE_DICT: Dictionary for all tensor blobs in the graph and their shapes
    shapes_wanted = []  # list of output names
    consts_wanted = []
    for op in OPS:
        for out in op.outputs:
            shape = out.get_shape()
            if not shape.is_fully_defined():
                shapes_wanted.append((compat.as_str_any(out.name), out))
            else:
                SHAPE_DICT[compat.as_str_any(out.name)] = shape.as_list()

        is_const = False
        if op.type == 'Const':
            is_const = True

        if op.type == 'Dequantize' and op.name in effectively_constant_ops:
            is_const = True

        if is_const:
            const = op.outputs[0]
            consts_wanted.append((compat.as_str_any(const.name), const))

    print('Collecting all the \'Const\' ops from the graph, by running it....')
    if len(shapes_wanted) > 0 or len(consts_wanted) > 0:
        tensor_names, tensors = zip(*(shapes_wanted + consts_wanted))
        if len(consts_wanted) > 0:
            const_tensor_names, _ = zip(*consts_wanted)
        else:
            const_tensor_names = []
        tensors_evaluated = sess.run(tensors, feed_dict=input_feed_dict)
        for i in range(len(tensor_names)):
            if tensor_names[i] not in SHAPE_DICT:
                SHAPE_DICT[tensor_names[i]] = list(tensors_evaluated[i].shape)
            if tensor_names[i] in const_tensor_names and tensor_names[
                    i] not in CONSTS:
                CONSTS[tensor_names[i]] = tensors_evaluated[i]
    print('Done.')

    # Fill in output information
    for op in OPS:
        output_names = set([compat.as_str_any(x.name) for x in op.outputs])
        if any(filter(output_names.__contains__, output_feature_names)):
            # retrieve model outputs
            for output in [
                    x for x in op.outputs if x.name in output_feature_names
            ]:
                #infer shape for Core ML
                tf_shape = SHAPE_DICT[compat.as_str_any(output.name)]
                shape = _infer_coreml_output_shape(tf_shape)
                out_name = output.name
                if shape is None:
                    output_features.append((compat.as_str_any(out_name), None))
                else:
                    output_features.append(
                        (compat.as_str_any(out_name), datatypes.Array(*shape)))

    if len(output_features) != len(output_feature_names):
        all_out_names_in_graph = [out_[0] for out_ in output_features]
        for given_out_name in output_feature_names:
            if given_out_name not in all_out_names_in_graph:
                raise ValueError(
                    "output name: {}, was provided, but the Tensorflow graph does not contain a tensor with this name."
                    .format(given_out_name))

    if not add_custom_layers:
        _check_unsupported_ops(OPS, output_feature_names,
                               effectively_constant_ops + unused_ops)
    print('Now starting translation to CoreML graph.')

    # Load all the dictionaries in the object of the class "context"
    context = Context(CONSTS, SHAPE_DICT, OPS, BLOB_GRAPH, output_features)

    # Interpret Input shapes and fill in input information for Core ML
    # (now that SHAPE_DICT and CONSTS are complete)
    sequence_inputs = dict()
    for input_tensor in input_feed_dict:
        input_name = compat.as_str_any(input_tensor.name)
        shape = SHAPE_DICT[input_name]

        if context.use_dfs_shape_infer:
            status = interpret_shape(input_name, context)
        else:
            status = False
        if status:
            print('Automatic shape interpretation succeeded for input blob %s' \
                %(input_name))
            shape = context.shape_dict_rank_4[input_name]

        if len(shape) == 4 and shape[0] != 1:
            sequence_inputs[input_name] = shape[0]

        # if the consumer of input_tensor is an one-hot encoding op,
        # treat it as a sequence.
        consumer_op = input_tensor.consumers()[0]
        if consumer_op.type == 'OneHot':
            shape = [
                1,
            ]
            sequence_inputs[input_name] = -1
        else:
            shape = _infer_coreml_input_shape(shape)
        input_features.append(
            (compat.as_str_any(input_name), datatypes.Array(*shape)))

    # Set classifier flag
    is_classifier = class_labels is not None
    mode = 'classifier' if is_classifier else None

    # Convert the TF graph with builder
    input_features = list(input_features)
    output_features = list(output_features)
    builder = NeuralNetworkBuilder(input_features, output_features, mode=mode)
    context.builder = builder
    context.session = sess
    context.input_feed_dict = input_feed_dict
    context.unused_ops = unused_ops
    context.effectively_constant_ops = effectively_constant_ops
    context.add_custom_layers = add_custom_layers
    context.custom_conversion_functions = custom_conversion_functions
    convert_ops_to_layers(context)
    sess.close()

    #optimizations on the nn spec
    optimize_nn_spec(spec=builder.spec)

    #Add a description for inputs that are sequences
    for i, inputs in enumerate(builder.spec.description.input):
        if inputs.name in sequence_inputs:
            seq_length = sequence_inputs[inputs.name]
            proto_shape = []
            if inputs.type.HasField('multiArrayType'):
                proto_shape = [
                    int(s) for s in inputs.type.multiArrayType.shape
                ]
            if seq_length == -1:
                msg = 'This input is a sequence'
                if len(proto_shape):
                    msg += '. Feed it an MLMultiArray of shape {} at runtime'.format(
                        str(['Seq_size', '1'] + proto_shape))
            else:
                msg = 'This input is a sequence of length ' + str(seq_length)
                if len(proto_shape):
                    msg += '. Feed it an MLMultiArray of shape {} at runtime'.format(
                        str([seq_length, 1] + proto_shape))
            builder.spec.description.input[i].shortDescription = msg

    # Add image input identifier
    if image_input_names is not None and isinstance(image_input_names,
                                                    _string_types):
        image_input_names = [image_input_names]

    # Replace all input/output blob names with ":" to "__" for compatible
    # auto-generated Objective C / Swift code
    interface_blob_names = []
    for idx, in_blob in enumerate(builder.spec.description.input):
        interface_blob_names.append(in_blob.name)
        builder.spec.description.input[idx].name = in_blob.name.replace(
            ':', '__').replace('/', '__')
    for idx, out_blob in enumerate(builder.spec.description.output):
        interface_blob_names.append(out_blob.name)
        builder.spec.description.output[idx].name = out_blob.name.replace(
            ':', '__').replace('/', '__')

    nn_spec = builder.nn_spec
    for i, spec_layer in enumerate(nn_spec.layers):
        for j, blob in enumerate(spec_layer.input):
            name = spec_layer.input[j]
            if name in interface_blob_names:
                spec_layer.input[j] = name.replace(':',
                                                   '__').replace('/', '__')
        for j, blob in enumerate(spec_layer.output):
            name = spec_layer.output[j]
            if name in interface_blob_names:
                spec_layer.output[j] = name.replace(':',
                                                    '__').replace('/', '__')

    if image_input_names is not None:
        for i, img in enumerate(image_input_names):
            image_input_names[i] = img.replace(':', '__').replace('/', '__')

    # Add classifier classes (if applicable)
    if is_classifier:
        classes_in = class_labels
        if isinstance(classes_in, _string_types):
            import os
            if not os.path.isfile(classes_in):
                raise ValueError("Path to class labels (%s) does not exist." % \
                    classes_in)
            with open(classes_in, 'r') as f:
                classes = f.read()
            classes = classes.splitlines()
        elif type(classes_in) is list:  # list[int or str]
            classes = classes_in
        else:
            raise ValueError('Class labels must be a list of integers / strings,'\
                ' or a file path')

        if predicted_feature_name is not None:
            builder.set_class_labels(
                classes,
                predicted_feature_name=predicted_feature_name,
                prediction_blob=predicted_probabilities_output)
        else:
            builder.set_class_labels(classes)

    # Set pre-processing parameters
    builder.set_pre_processing_parameters(image_input_names=image_input_names,
                                          is_bgr=is_bgr,
                                          red_bias=red_bias,
                                          green_bias=green_bias,
                                          blue_bias=blue_bias,
                                          gray_bias=gray_bias,
                                          image_scale=image_scale)

    print(
        "Translation to CoreML spec completed. Now compiling and saving the CoreML model."
    )
    try:
        import coremltools
        coremltools.models.utils.save_spec(builder.spec, mlmodel_path)
        mlmodel = MLModel(builder.spec)
    except RuntimeError as e:
        raise ValueError('Compilation failed: {}'.format(str(e)))

    print("\n Core ML model generated. Saved at location: %s \n" %
          (mlmodel_path))
    print('Core ML input(s): \n', builder.spec.description.input)
    print('Core ML output(s): \n', builder.spec.description.output)

    # print information about all ops for which custom layers have been added
    if len(context.ops_converted_to_custom_layers) > 0:
        print('\n')
        print("Custom layers have been added to the CoreML model "
              "corresponding to the following ops in the TF graph: ")
        for i, op in enumerate(context.ops_converted_to_custom_layers):
            input_info = []
            for input_ in op.inputs:
                input_info.append(
                    (str(input_.name),
                     context.shape_dict.get(input_.name,
                                            str("Shape not available"))))
            output_info = []
            for output_ in op.outputs:
                output_info.append(
                    (str(output_.name),
                     context.shape_dict.get(output_.name,
                                            str("Shape not available"))))
            print(
                "{}/{}: op type: {}, op input names and shapes: {}, op output names and shapes: {}"
                .format(i + 1, len(context.ops_converted_to_custom_layers),
                        op.type, str(input_info), str(output_info)))

    # Return the protobuf model
    return mlmodel
Exemplo n.º 5
0
def convert(model,
            mode=None,
            image_input_names=[],
            preprocessing_args={},
            image_output_names=[],
            deprocessing_args={},
            class_labels=None,
            predicted_feature_name='classLabel'):
    """
    Convert ONNX model to CoreML.
    Parameters
    ----------
    model: ONNX model | str
        An ONNX model with parameters loaded in onnx package or path to file
        with models.
    mode: str ('classifier', 'regressor' or None)
        Mode of the converted coreml model:
        'classifier', a NeuralNetworkClassifier spec will be constructed.
        'regressor', a NeuralNetworkRegressor spec will be constructed.
    preprocessing_args: dict
        'is_bgr', 'red_bias', 'green_bias', 'blue_bias', 'gray_bias',
        'image_scale' keys with the same meaning as
        https://apple.github.io/coremltools/generated/coremltools.models.neural_network.html#coremltools.models.neural_network.NeuralNetworkBuilder.set_pre_processing_parameters
    deprocessing_args: dict
        Same as 'preprocessing_args' but for deprocessing.
    class_labels: A string or list of strings.
        As a string it represents the name of the file which contains
        the classification labels (one per line).
        As a list of strings it represents a list of categories that map
        the index of the output of a neural network to labels in a classifier.
    predicted_feature_name: str
        Name of the output feature for the class labels exposed in the Core ML
        model (applies to classifiers only). Defaults to 'classLabel'
    Returns
    -------
    model: A coreml model.
    """
    if isinstance(model, basestring):
        onnx_model = onnx.load(model)
    elif isinstance(model, onnx.ModelProto):
        onnx_model = model
    else:
        raise TypeError(
            "Model must be file path to .onnx file or onnx loaded model"
        )

    transformers = [
        ReshapeInitTensorFuser(),
        DropoutRemover(),
        ConvAddFuser(),
        BNBroadcastedMulFuser(),
        BNBroadcastedAddFuser(),
        PixelShuffleFuser(),
        DanglingOutputsRemover()
    ]

    graph = _prepare_onnx_graph(onnx_model.graph, transformers)

    input_features = _features(graph.inputs)
    output_features = _features(graph.outputs, adapt_shape=False)

    is_deprocess_bgr_only = (len(deprocessing_args) == 1) and \
                            ("is_bgr" in deprocessing_args)
    add_deprocess = (len(image_output_names) > 0) and \
                    (len(deprocessing_args) > 0) and \
                    (not is_deprocess_bgr_only)

    if add_deprocess:
        mapping = {}
        for f in output_features:
            output_name = f[0]
            mapping[output_name] = graph.get_unique_edge_name(output_name)
        graph = OutputRenamer(mapping)(graph)

    builder = NeuralNetworkBuilder(input_features, output_features, mode)

    if len(image_input_names) > 0:
        builder.set_pre_processing_parameters(
            image_input_names=image_input_names,
            is_bgr=preprocessing_args.get('is_bgr', False),
            red_bias=preprocessing_args.get('red_bias', 0.0),
            green_bias=preprocessing_args.get('green_bias', 0.0),
            blue_bias=preprocessing_args.get('blue_bias', 0.0),
            gray_bias=preprocessing_args.get('gray_bias', 0.0),
            image_scale=preprocessing_args.get('image_scale', 1.0)
        )

    if len(image_output_names) > 0:
        for f in output_features:
            f_name = f[0]
            if f_name in image_output_names:
                is_bgr = deprocessing_args.get('is_bgr', False)
                _convert_multiarray_output_to_image(
                    builder.spec, f_name, is_bgr=is_bgr
                )

    for node in graph.nodes:
        _convert_node(builder, node)

    if add_deprocess:
        for f in output_features:
            output_name = f[0]
            if output_name not in image_output_names:
                continue
            output_shape = f[1].dimensions
            if len(output_shape) == 2 or output_shape[0] == 1:
                is_grayscale = True
            elif output_shape[0] == 3:
                is_grayscale = False
            else:
                raise ValueError('Output must be RGB image or Grayscale')
            _set_deprocessing(
                is_grayscale,
                builder,
                deprocessing_args,
                mapping[output_name],
                output_name
            )

    if class_labels is not None:
        if type(class_labels) is str:
            labels = [l.strip() for l in open(class_labels).readlines()]
        elif type(class_labels) is list:
            labels = class_labels
        else:
            raise TypeError(
                "synset variable of unknown type. Type found: {}. \
                Expected either string or list of strings."
                .format(type(class_labels),))

        builder.set_class_labels(
            class_labels=labels,
            predicted_feature_name=predicted_feature_name
        )

    return MLModel(builder.spec)
Exemplo n.º 6
0
def convert(
        model,  # type: Union[onnx.ModelProto, Text]
        mode=None,  # type: Optional[Text]
        image_input_names=[],  # type: Sequence[Text]
        preprocessing_args={},  # type: Dict[Text, Any]
        image_output_names=[],  # type: Sequence[Text]
        deprocessing_args={},  # type: Dict[Text, Any]
        class_labels=None,  # type: Union[Text, Iterable[Text], None]
        predicted_feature_name='classLabel',  # type: Text
        add_custom_layers=False,  # type: bool
        custom_conversion_functions={},  #type: Dict[Text, Any]
        onnx_coreml_input_shape_map={},  # type: Dict[Text, List[int,...]]
        disable_coreml_rank5_mapping=False):
    # type: (...) -> MLModel
    """
    Convert ONNX model to CoreML.
    Parameters
    ----------
    model:
        An ONNX model with parameters loaded in onnx package or path to file
        with models.
    mode: 'classifier', 'regressor' or None
        Mode of the converted coreml model:
        'classifier', a NeuralNetworkClassifier spec will be constructed.
        'regressor', a NeuralNetworkRegressor spec will be constructed.
    preprocessing_args:
        'is_bgr', 'red_bias', 'green_bias', 'blue_bias', 'gray_bias',
        'image_scale' keys with the same meaning as
        https://apple.github.io/coremltools/generated/coremltools.models.neural_network.html#coremltools.models.neural_network.NeuralNetworkBuilder.set_pre_processing_parameters
    deprocessing_args:
        Same as 'preprocessing_args' but for deprocessing.
    class_labels:
        As a string it represents the name of the file which contains
        the classification labels (one per line).
        As a list of strings it represents a list of categories that map
        the index of the output of a neural network to labels in a classifier.
    predicted_feature_name:
        Name of the output feature for the class labels exposed in the Core ML
        model (applies to classifiers only). Defaults to 'classLabel'
    add_custom_layers: bool
        Flag to turn on addition of custom CoreML layers for unsupported ONNX ops or attributes within
        a supported op.
    custom_conversion_functions: dict()
        A dictionary with keys corresponding to the names/types of onnx ops and values as functions taking
        an object of class coreml-tools's 'NeuralNetworkBuilder', Graph' (see onnx-coreml/_graph.Graph),
        'Node' (see onnx-coreml/_graph.Node), ErrorHandling (see onnx-coreml/_error_utils.ErrorHandling).
        This custom conversion function gets full control and responsibility for converting given onnx op.
        This function returns nothing and is responsible for adding a equivalent CoreML layer via 'NeuralNetworkBuilder'
    onnx_coreml_input_shape_map: dict()
        (Optional) A dictionary with keys corresponding to the model input names. Values are a list of integers that specify
        how the shape of the input is mapped to CoreML. Convention used for CoreML shapes is
        0: Sequence, 1: Batch, 2: channel, 3: height, 4: width.
        For example, an input of rank 2 could be mapped as [3,4] (i.e. H,W) or [1,2] (i.e. B,C) etc.
        This is ignored if "disable_coreml_rank5_mapping" is set to True.
    disable_coreml_rank5_mapping: bool
        If True, then it disables the "RANK5_ARRAY_MAPPING" or enables the "EXACT_ARRAY_MAPPING"
        option in CoreML (https://github.com/apple/coremltools/blob/655b3be5cc0d42c3c4fa49f0f0e4a93a26b3e492/mlmodel/format/NeuralNetwork.proto#L67)
        Thus, no longer, onnx tensors are forced to map to rank 5 CoreML tensors.
        With this flag on, a rank r ONNX tensor, (1<=r<=5), will map to a rank r tensor in CoreML as well.
        This flag must be on to utilize any of the new layers added in CoreML 3 (i.e. specification version 4, iOS13)

    Returns
    -------
    model: A coreml model.
    """
    if isinstance(model, Text):
        onnx_model = onnx.load(model)
    elif isinstance(model, onnx.ModelProto):
        onnx_model = model
    else:
        raise TypeError(
            "Model must be file path to .onnx file or onnx loaded model")

    global USE_SHAPE_MAPPING
    if disable_coreml_rank5_mapping:
        USE_SHAPE_MAPPING = False
    '''
    First, apply a few optimizations to the ONNX graph,
    in preparation for conversion to CoreML. 
    '''

    # Using Dummy transformation to conditionally disable certain transformation
    class DummyTransformation(object):
        def __call__(self, graph):
            return graph

    transformers = [
        ConstantsToInitializers(),
        ShapeOpRemover(),
        ReshapeInitTensorFuser(),
        DropoutRemover(),
        UnsqueezeConstantRemover(),
        TransposeConstantRemover(),
        SliceConstantRemover(),
        ConcatConstantRemover(),
        ConvAddFuser(),
        BNBroadcastedMulFuser(),
        BNBroadcastedAddFuser(),
        ReshapeTransposeReshape_pattern1(),
        PixelShuffleFuser(),
        AddModelInputsOutputs()
        if not disable_coreml_rank5_mapping else DummyTransformation(),
        DivMulConstantRemover(),
        GatherConstantRemover(),
        ConstantFillToInitializers(),
    ]  # type: Iterable[Transformer]

    onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
    graph = _prepare_onnx_graph(onnx_model.graph, transformers)
    '''
    Check for ImageScalar nodes in ONNX, this will indicate whether input image preprocessing needs
    to be added to the CoreML graph or not. 
    '''
    # are there ImageScaler nodes in the Graph?
    # If yes then add the info from it to the "preprocessing_args" dictionary, if the dictionary is not
    # already provided by the user
    if not bool(preprocessing_args):
        for node in graph.nodes:
            if node.op_type == 'ImageScaler':
                inp_name = node.inputs[0]
                scale = node.attrs.get('scale', 1.0)
                bias = node.attrs.get('bias', [0, 0, 0])
                if not (len(bias) == 1 or len(bias) == 3):
                    continue
                if 'image_scale' in preprocessing_args:
                    preprocessing_args['image_scale'][inp_name] = scale
                else:
                    preprocessing_args['image_scale'] = {inp_name: scale}
                if len(bias) == 3:
                    for i, color in enumerate(['red', 'green', 'blue']):
                        if color + '_bias' in preprocessing_args:
                            preprocessing_args[color +
                                               '_bias'][inp_name] = bias[i]
                        else:
                            preprocessing_args[color + '_bias'] = {
                                inp_name: bias[i]
                            }
                else:
                    if 'gray_bias' in preprocessing_args:
                        preprocessing_args['gray_bias'][inp_name] = bias[0]
                    else:
                        preprocessing_args['gray_bias'] = {inp_name: bias[0]}
                if inp_name not in image_input_names:
                    image_input_names.append(inp_name)  # type: ignore

    # remove all ImageScaler ops
    graph = graph.transformed([ImageScalerRemover()])
    '''
    Gather information (name, shape) for model inputs and outputs
    This information is then used to initialize the neural network builder object of coremltools. 
    The builder object is later used to add layers to the CoreML model. 
    '''

    #Make CoreML input and output features by gathering shape info and
    #interpreting it for CoreML
    input_features = _make_coreml_input_features(graph,
                                                 onnx_coreml_input_shape_map,
                                                 disable_coreml_rank5_mapping)
    if len(image_output_names) > 0:
        output_features = _make_coreml_output_features(
            graph,
            forceShape=True,
            disable_coreml_rank5_mapping=disable_coreml_rank5_mapping)
    else:
        output_features = _make_coreml_output_features(
            graph, disable_coreml_rank5_mapping=disable_coreml_rank5_mapping)

    builder = NeuralNetworkBuilder(
        input_features,
        output_features,
        mode=mode,
        disable_rank5_shape_mapping=disable_coreml_rank5_mapping)
    '''
    Set CoreML input,output types (float, double, int) same as onnx types, if supported
    '''
    _transform_coreml_dtypes(builder, graph.inputs, graph.outputs)
    '''what follows is some book-keeping to support outputs of type image. 
    '''

    is_deprocess_bgr_only = (len(deprocessing_args) == 1) and \
                            ("is_bgr" in deprocessing_args)
    add_deprocess = (len(image_output_names) > 0) and \
                    (len(deprocessing_args) > 0) and \
                    (not is_deprocess_bgr_only)

    if add_deprocess:
        mapping = {}
        for f in output_features:
            output_name = f[0]
            mapping[output_name] = graph.get_unique_edge_name(output_name)
        graph = OutputRenamer(mapping)(graph)

    if len(image_input_names) > 0:
        builder.set_pre_processing_parameters(
            image_input_names=image_input_names,
            is_bgr=preprocessing_args.get('is_bgr', False),
            red_bias=preprocessing_args.get('red_bias', 0.0),
            green_bias=preprocessing_args.get('green_bias', 0.0),
            blue_bias=preprocessing_args.get('blue_bias', 0.0),
            gray_bias=preprocessing_args.get('gray_bias', 0.0),
            image_scale=preprocessing_args.get('image_scale', 1.0))

    preprocessing_args.clear()

    if len(image_output_names) > 0:
        for f in output_features:
            f_name = f[0]
            if f_name in image_output_names:
                is_bgr = deprocessing_args.get('is_bgr', False)
                _convert_multiarray_output_to_image(builder.spec,
                                                    f_name,
                                                    is_bgr=is_bgr)
    '''
    Iterate through all the ONNX ops and translate them to CoreML layers, one by one. 
    '''
    '''
    before proceeding to start the layer translation process,
    check whether there is an op in the ONNX graph, whose translation function is not yet
    implemented in the converter or which is not supported in the CoreML framework. If so, 
    raise an error before starting the process.
    (if the user desires to add a custom layer then this check is not required)
    '''
    if not add_custom_layers:
        _check_unsupported_ops(graph.nodes, disable_coreml_rank5_mapping)
    '''
    ErrorHandling is a generic class, useful to store a variety of parameters during the conversion process  
    '''
    err = ErrorHandling(
        add_custom_layers,
        custom_conversion_functions,
        disable_coreml_rank5_mapping=disable_coreml_rank5_mapping)

    for i, node in enumerate(graph.nodes):
        print("%d/%d: Converting Node Type %s" %
              (i + 1, len(graph.nodes), node.op_type))
        if disable_coreml_rank5_mapping:
            _convert_node_nd(builder, node, graph, err)
        else:
            _add_const_inputs_if_required(builder, node, graph, err)
            _convert_node(builder, node, graph, err)

    if DEBUG:
        plot_graph(graph,
                   graph_img_path='/tmp/after_conversion.pdf',
                   show_coreml_mapped_shapes=not disable_coreml_rank5_mapping)

    if add_deprocess:
        for f in output_features:
            output_name = f[0]
            if output_name not in image_output_names:
                continue
            output_shape = f[1].dimensions
            if len(output_shape) == 2 or output_shape[0] == 1:
                is_grayscale = True
            elif output_shape[0] == 3:
                is_grayscale = False
            else:
                raise ValueError('Output must be RGB image or Grayscale')
            _set_deprocessing(is_grayscale, builder, deprocessing_args,
                              mapping[output_name], output_name)

    if class_labels is not None:
        if isinstance(class_labels, Text):
            labels = [l.strip() for l in open(class_labels).readlines()
                      ]  # type: Sequence[Text]
        elif isinstance(class_labels, list):
            labels = class_labels
        else:
            raise TypeError("synset variable of unknown type. Type found: {}. \
                Expected either string or list of strings.".format(
                type(class_labels), ))

        builder.set_class_labels(class_labels=labels,
                                 predicted_feature_name=predicted_feature_name)

    def _add_informative_description(feature, raise_error=True):
        if feature.type.WhichOneof('Type') == 'multiArrayType':
            if feature.name in graph.onnx_coreml_shape_mapping and feature.name in graph.shape_dict:
                mapp = graph.onnx_coreml_shape_mapping[feature.name]
                onnx_shape = graph.shape_dict[feature.name]
                if raise_error:
                    assert len(mapp) == len(
                        onnx_shape), "Something wrong in shape"
                if len(mapp) == len(onnx_shape):
                    shape = []
                    for i in range(5):
                        if i in mapp:
                            shape += [int(onnx_shape[mapp.index(i)])]
                        else:
                            shape += [1]
                    msg = 'MultiArray of shape {}. The first and second dimensions correspond to sequence and batch size, respectively'.format(
                        str(tuple(shape)))
                    feature.shortDescription += msg

    optional_input_names = []
    for tup in graph.optional_inputs:
        optional_input_names.append(tup[0])
    optional_output_names = []
    for tup in graph.optional_outputs:
        optional_output_names.append(tup[0])

    # add description for inputs and outputs shapes
    remove_input_id = []
    for i, input_ in enumerate(builder.spec.description.input):
        if input_.name not in optional_input_names:
            if not disable_coreml_rank5_mapping:
                _add_informative_description(input_)
        else:
            remove_input_id.append(i)
    remove_output_id = []
    for i, output_ in enumerate(builder.spec.description.output):
        if output_.name not in optional_output_names:
            if not disable_coreml_rank5_mapping:
                _add_informative_description(output_, raise_error=False)
        else:
            remove_output_id.append(i)

    for index in sorted(remove_input_id, reverse=True):
        del builder.spec.description.input[index]
    for index in sorted(remove_output_id, reverse=True):
        del builder.spec.description.output[index]

    if len(graph.optional_inputs) > 0 or len(graph.optional_outputs):
        builder.add_optionals(graph.optional_inputs, graph.optional_outputs)

    print(
        "Translation to CoreML spec completed. Now compiling the CoreML model."
    )
    try:
        if DEBUG:
            import coremltools
            coremltools.models.utils.save_spec(
                builder.spec, '/tmp/node_model_raw_spec.mlmodel')
            from coremltools.models.neural_network.printer import print_network_spec
            print_network_spec(builder.spec, style='coding')
        mlmodel = MLModel(builder.spec)
    except RuntimeError as e:
        raise ValueError('Compilation failed: {}'.format(str(e)))
    print('Model Compilation done.')

    # print information about all ops for which custom layers have been added
    if len(err.custom_layer_nodes) > 0:
        print('\n')
        print("Custom layers have been added to the CoreML model "
              "corresponding to the following ops in the onnx model: ")
        for i, node in enumerate(err.custom_layer_nodes):
            input_info = []
            for input_ in node.inputs:
                input_info.append(
                    (str(input_),
                     graph.shape_dict.get(input_, str("Shape not available"))))
            output_info = []
            for output_ in node.outputs:
                output_info.append(
                    (str(output_),
                     graph.shape_dict.get(output_,
                                          str("Shape not available"))))
            print(
                "{}/{}: op type: {}, op input names and shapes: {}, op output names and shapes: {}"
                .format(i + 1, len(err.custom_layer_nodes), node.op_type,
                        str(input_info), str(output_info)))

    return mlmodel
Exemplo n.º 7
0
def _convert_pb_to_mlmodel(tf_model_path,
                           mlmodel_path,
                           output_feature_names,
                           input_name_shape_dict=None,
                           image_input_names=None,
                           is_bgr=False,
                           red_bias=0.0,
                           green_bias=0.0,
                           blue_bias=0.0,
                           gray_bias=0.0,
                           image_scale=1.0,
                           class_labels=None,
                           predicted_feature_name=None,
                           predicted_probabilities_output=''):
    if input_name_shape_dict is None:
        input_name_shape_dict = {}
    # Load the TF graph
    with open(tf_model_path, 'rb') as f:
        serialized = f.read()

    tf.reset_default_graph()
    gdef = tf.GraphDef()
    gdef.ParseFromString(serialized)

    with tf.Graph().as_default() as g:
        tf.import_graph_def(gdef, name='')

    sess = tf.Session(graph=g)
    OPS = g.get_operations()
    OPS = _topological_sort_ops(OPS)
    _check_unsupported_ops(OPS, output_feature_names)

    SHAPE_DICT = {}  #Tensor name --> shape ({str: list})
    CONSTS = {}  #Const Tensor name --> value
    BLOB_GRAPH = {}  #Blob name to list of ops it feeds into

    # Make Dictionary of Input blob to the list of ops it feeds into
    for op in OPS:
        for inp in op.inputs:
            if inp.name in BLOB_GRAPH:
                BLOB_GRAPH[inp.name].append(op)
        for out in op.outputs:
            if out.name not in BLOB_GRAPH:
                BLOB_GRAPH[out.name] = []

    # Fill in input information
    input_features = []
    output_features = []
    input_feed_dict = dict()  #Input tensors' values

    # run through all placeholders
    for op in OPS:
        output_names = set([compat.as_bytes(x.name) for x in op.outputs])
        if op.type == 'Placeholder':
            # Handle placeholders -- all placeholders are inputs
            assert not filter(output_names.__contains__, output_feature_names), \
                ('Output feature cannot be a placeholder')
            input_name = compat.as_bytes(op.outputs[0].name)
            shape = op.outputs[0].get_shape()
            if not (shape.is_fully_defined()
                    or input_name in input_name_shape_dict):
                assert False, ("%s is a placehoder with incomplete shape %s" %
                               (input_name, str(shape)))
            if shape.is_fully_defined():
                shape = shape.as_list()
            else:
                shape = input_name_shape_dict[input_name]

            if len(shape) == 0:  # scalar - use a 1
                input_feed_dict[op.outputs[0]] = 1
            else:
                input_feed_dict[op.outputs[0]] = np.random.rand(*shape)

            SHAPE_DICT[input_name] = shape

    # Populate SHAPE_DICT:
    # Dictionary for all tensor blobs in the graph and their shapes
    shapes_wanted = []
    for op in OPS:
        for out in op.outputs:
            shape = out.get_shape()
            if not shape.is_fully_defined():
                shapes_wanted.append((compat.as_bytes(out.name), out))
            else:
                SHAPE_DICT[compat.as_bytes(out.name)] = shape.as_list()

    if len(shapes_wanted) > 0:
        print("Shapes not found for %d tensors. "
              "Executing graph to determine shapes. " % (len(shapes_wanted)))
        tensor_names, tensors = zip(*shapes_wanted)
        tensors_evaluated = sess.run(tensors, feed_dict=input_feed_dict)
        for i in range(len(tensor_names)):
            SHAPE_DICT[tensor_names[i]] = list(tensors_evaluated[i].shape)

    # Fill in output information and CONSTS dictionary
    for op in OPS:
        output_names = set([compat.as_bytes(x.name) for x in op.outputs])
        if filter(output_names.__contains__, output_feature_names):
            # retrieve model outputs
            for output in [
                    x for x in op.outputs if x.name in output_feature_names
            ]:
                #infer shape for Core ML
                tf_shape = SHAPE_DICT[compat.as_bytes(output.name)]
                shape = _infer_coreml_output_shape(tf_shape)
                out_name = output.name
                if shape is None:
                    output_features.append((compat.as_bytes(out_name), None))
                else:
                    output_features.append(
                        (compat.as_bytes(out_name), datatypes.Array(*shape)))
        elif op.type == 'Const':
            # retrieve all consts and store them in dictionary
            const = op.outputs[0]
            CONSTS[compat.as_bytes(const.name)] = sess.run(
                const, feed_dict=input_feed_dict)

    assert len(output_features) == len(output_feature_names), (
        'Tensorflow Graph does not contain all the provided Output name(s)')

    # Load all the dictionaries in the object of class context
    context = Context(CONSTS, SHAPE_DICT, OPS, BLOB_GRAPH, output_features)

    # Interpret Input shapes and fill in input information for Core ML
    # (now that SHAPE_DICT and CONSTS are complete)
    sequence_inputs = dict()
    for input_tensor in input_feed_dict:
        input_name = compat.as_bytes(input_tensor.name)
        shape = SHAPE_DICT[input_name]

        if context.use_dfs_shape_infer:
            status = interpret_shape(input_name, context)
        else:
            status = False
        if status:
            print('Automatic shape interpretation succeeded for input blob %s' \
                %(input_name))
            shape = context.shape_dict_rank_4[input_name]

        if len(shape) == 4 and shape[0] != 1:
            sequence_inputs[input_name] = shape[0]

        # if the consumer of input_tensor is an one-hot encoding op,
        # treat it as a sequence.
        consumer_op = input_tensor.consumers()[0]
        if consumer_op.type == 'OneHot':
            shape = [
                1,
            ]
            sequence_inputs[input_name] = -1
        else:
            shape = _infer_coreml_input_shape(shape)
        input_features.append(
            (compat.as_bytes(input_name), datatypes.Array(*shape)))

    # Set classifier flag
    is_classifier = class_labels is not None
    mode = 'classifier' if is_classifier else None

    # Convert the TF graph with builder
    input_features = list(input_features)
    output_features = list(output_features)
    builder = NeuralNetworkBuilder(input_features, output_features, mode=mode)
    context.builder = builder
    context.session = sess
    context.input_feed_dict = input_feed_dict
    convert_ops_to_layers(context)
    sess.close()

    #optimizations on the nn spec
    optimize_nn_spec(builder=builder)

    #Add a description for inputs that are sequences
    for i, inputs in enumerate(builder.spec.description.input):
        if inputs.name in sequence_inputs:
            seq_length = sequence_inputs[inputs.name]
            if seq_length == -1:
                builder.spec.description.input[i].shortDescription = \
                  'This input is a sequence'
            else:
                builder.spec.description.input[i].shortDescription = \
                  'This input is a sequence of length ' + str(seq_length)

    # Add image input identifier
    if image_input_names is not None and isinstance(image_input_names,
                                                    _string_types):
        image_input_names = [image_input_names]

    # Replace all input/output blob names with ":" to "__" for compatible
    # auto-generated Objective C / Swift code
    interface_blob_names = []
    for idx, in_blob in enumerate(builder.spec.description.input):
        interface_blob_names.append(in_blob.name)
        builder.spec.description.input[idx].name = in_blob.name.replace(
            ':', '__').replace('/', '__')
    for idx, out_blob in enumerate(builder.spec.description.output):
        interface_blob_names.append(out_blob.name)
        builder.spec.description.output[idx].name = out_blob.name.replace(
            ':', '__').replace('/', '__')

    nn_spec = builder.nn_spec
    for i, spec_layer in enumerate(nn_spec.layers):
        for j, blob in enumerate(spec_layer.input):
            name = spec_layer.input[j]
            if name in interface_blob_names:
                spec_layer.input[j] = name.replace(':',
                                                   '__').replace('/', '__')
        for j, blob in enumerate(spec_layer.output):
            name = spec_layer.output[j]
            if name in interface_blob_names:
                spec_layer.output[j] = name.replace(':',
                                                    '__').replace('/', '__')

    if image_input_names is not None:
        for i, img in enumerate(image_input_names):
            image_input_names[i] = img.replace(':', '__').replace('/', '__')

    # Add classifier classes (if applicable)
    if is_classifier:
        classes_in = class_labels
        if isinstance(classes_in, _string_types):
            import os
            if not os.path.isfile(classes_in):
                raise ValueError("Path to class labels (%s) does not exist." % \
                    classes_in)
            with open(classes_in, 'r') as f:
                classes = f.read()
            classes = classes.splitlines()
        elif type(classes_in) is list:  # list[int or str]
            classes = classes_in
        else:
            raise ValueError('Class labels must be a list of integers / strings,'\
                ' or a file path')

        if predicted_feature_name is not None:
            builder.set_class_labels(
                classes,
                predicted_feature_name=predicted_feature_name,
                prediction_blob=predicted_probabilities_output)
        else:
            builder.set_class_labels(classes)

    # Set pre-processing paramsters
    builder.set_pre_processing_parameters(image_input_names=image_input_names,
                                          is_bgr=is_bgr,
                                          red_bias=red_bias,
                                          green_bias=green_bias,
                                          blue_bias=blue_bias,
                                          gray_bias=gray_bias,
                                          image_scale=image_scale)

    utils.save_spec(builder.spec, mlmodel_path)
    print("\n Core ML model generated. Saved at location: %s \n" %
          (mlmodel_path))
    print('Core ML input(s): \n', builder.spec.description.input)
    print('Core ML output(s): \n', builder.spec.description.output)

    # Return the protobuf spec
    spec = builder.spec
    return MLModel(spec)
Exemplo n.º 8
0
def convert(
        model,  # type: Union[onnx.ModelProto, Text]
        mode=None,  # type: Optional[Text]
        image_input_names=[],  # type: Sequence[Text]
        preprocessing_args={},  # type: Dict[Text, Any]
        image_output_names=[],  # type: Sequence[Text]
        deprocessing_args={},  # type: Dict[Text, Any]
        class_labels=None,  # type: Union[Text, Iterable[Text], None]
        predicted_feature_name='classLabel',  # type: Text
        add_custom_layers=False,  # type: bool
        custom_conversion_functions={},  #type: Dict[Text, Any]
):
    # type: (...) -> MLModel
    """
    Convert ONNX model to CoreML.
    Parameters
    ----------
    model:
        An ONNX model with parameters loaded in onnx package or path to file
        with models.
    mode: 'classifier', 'regressor' or None
        Mode of the converted coreml model:
        'classifier', a NeuralNetworkClassifier spec will be constructed.
        'regressor', a NeuralNetworkRegressor spec will be constructed.
    preprocessing_args:
        'is_bgr', 'red_bias', 'green_bias', 'blue_bias', 'gray_bias',
        'image_scale' keys with the same meaning as
        https://apple.github.io/coremltools/generated/coremltools.models.neural_network.html#coremltools.models.neural_network.NeuralNetworkBuilder.set_pre_processing_parameters
    deprocessing_args:
        Same as 'preprocessing_args' but for deprocessing.
    class_labels:
        As a string it represents the name of the file which contains
        the classification labels (one per line).
        As a list of strings it represents a list of categories that map
        the index of the output of a neural network to labels in a classifier.
    predicted_feature_name:
        Name of the output feature for the class labels exposed in the Core ML
        model (applies to classifiers only). Defaults to 'classLabel'
    Returns
    -------
    model: A coreml model.
    """
    if isinstance(model, Text):
        onnx_model = onnx.load(model)
    elif isinstance(model, onnx.ModelProto):
        onnx_model = model
    else:
        raise TypeError(
            "Model must be file path to .onnx file or onnx loaded model")

    transformers = [
        ConstantsToInitializers(),
        ReshapeInitTensorFuser(),
        DropoutRemover(),
        ConvAddFuser(),
        BNBroadcastedMulFuser(),
        BNBroadcastedAddFuser(),
        PixelShuffleFuser(),
        AddModelInputsOutputs(),
    ]  # type: Iterable[Transformer]

    graph = _prepare_onnx_graph(onnx_model.graph, transformers)

    #Make CoreML input and output features by gathering shape info and
    #interpreting it for CoreML
    input_features = _make_coreml_input_features(graph)
    output_features = _make_coreml_output_features(graph)

    builder = NeuralNetworkBuilder(input_features, output_features)
    _transform_coreml_dtypes(builder, graph.inputs, graph.outputs)

    is_deprocess_bgr_only = (len(deprocessing_args) == 1) and \
                            ("is_bgr" in deprocessing_args)
    add_deprocess = (len(image_output_names) > 0) and \
                    (len(deprocessing_args) > 0) and \
                    (not is_deprocess_bgr_only)

    if add_deprocess:
        mapping = {}
        for f in output_features:
            output_name = f[0]
            mapping[output_name] = graph.get_unique_edge_name(output_name)
        graph = OutputRenamer(mapping)(graph)

    if len(image_input_names) > 0:
        builder.set_pre_processing_parameters(
            image_input_names=image_input_names,
            is_bgr=preprocessing_args.get('is_bgr', False),
            red_bias=preprocessing_args.get('red_bias', 0.0),
            green_bias=preprocessing_args.get('green_bias', 0.0),
            blue_bias=preprocessing_args.get('blue_bias', 0.0),
            gray_bias=preprocessing_args.get('gray_bias', 0.0),
            image_scale=preprocessing_args.get('image_scale', 1.0))

    if len(image_output_names) > 0:
        for f in output_features:
            f_name = f[0]
            if f_name in image_output_names:
                is_bgr = deprocessing_args.get('is_bgr', False)
                _convert_multiarray_output_to_image(builder.spec,
                                                    f_name,
                                                    is_bgr=is_bgr)
    '''Iterate through all the ops and translate them to CoreML layers.
    '''
    if not add_custom_layers:
        _check_unsupported_ops(graph.nodes)
    err = ErrorHandling(add_custom_layers, custom_conversion_functions)
    for i, node in enumerate(graph.nodes):
        print("%d/%d: Converting Node Type %s" %
              (i + 1, len(graph.nodes), node.op_type))
        _convert_node(builder, node, graph, err)

    if add_deprocess:
        for f in output_features:
            output_name = f[0]
            if output_name not in image_output_names:
                continue
            output_shape = f[1].dimensions
            if len(output_shape) == 2 or output_shape[0] == 1:
                is_grayscale = True
            elif output_shape[0] == 3:
                is_grayscale = False
            else:
                raise ValueError('Output must be RGB image or Grayscale')
            _set_deprocessing(is_grayscale, builder, deprocessing_args,
                              mapping[output_name], output_name)

    if class_labels is not None:
        if isinstance(class_labels, Text):
            labels = [l.strip() for l in open(class_labels).readlines()
                      ]  # type: Sequence[Text]
        elif isinstance(class_labels, list):
            labels = class_labels
        else:
            raise TypeError("synset variable of unknown type. Type found: {}. \
                Expected either string or list of strings.".format(
                type(class_labels), ))

        builder.set_class_labels(class_labels=labels,
                                 predicted_feature_name=predicted_feature_name)

    # add description to inputs/outputs that feed in/out of recurrent layers
    for node_ in graph.nodes:
        if str(node_.op_type) in _SEQUENCE_LAYERS_REGISTRY:
            input_ = node_.inputs[0]
            output_ = node_.outputs[0]
            for i, inputs in enumerate(builder.spec.description.input):
                if inputs.name == input_:
                    builder.spec.description.input[
                        i].shortDescription = 'This input is a sequence'
            for i, outputs in enumerate(builder.spec.description.output):
                if outputs.name == output_:
                    builder.spec.description.output[
                        i].shortDescription = 'This output is a sequence'

    mlmodel = MLModel(builder.spec)

    # print information about all ops for which custom layers have been added
    if len(err.custom_layer_nodes) > 0:
        print('\n')
        print("Custom layers have been added to the CoreML model "
              "corresponding to the following ops in the onnx model: ")
        for i, node in enumerate(err.custom_layer_nodes):
            input_info = []
            for input_ in node.inputs:
                input_info.append(
                    (str(input_),
                     graph.shape_dict.get(input_, str("Shape not available"))))
            output_info = []
            for output_ in node.outputs:
                output_info.append(
                    (str(output_),
                     graph.shape_dict.get(output_,
                                          str("Shape not available"))))
            print(
                "{}/{}: op type: {}, op input names and shapes: {}, op output names and shapes: {}"
                .format(i + 1, len(err.custom_layer_nodes), node.op_type,
                        str(input_info), str(output_info)))

    return mlmodel
Exemplo n.º 9
0
def convert(
    model,  # type: Union[onnx.ModelProto, Text]
    mode=None,  # type: Optional[Text]
    image_input_names=[],  # type: Sequence[Text]
    preprocessing_args={},  # type: Dict[Text, Any]
    image_output_names=[],  # type: Sequence[Text]
    deprocessing_args={},  # type: Dict[Text, Any]
    class_labels=None,  # type: Union[Text, Iterable[Text], None]
    predicted_feature_name='classLabel',  # type: Text
    add_custom_layers=False,  # type: bool
    custom_conversion_functions={},  #type: Dict[Text, Any]
    onnx_coreml_input_shape_map={}  # type: Dict[Text, List[int,...]]
):
    # type: (...) -> MLModel
    """
    Convert ONNX model to CoreML.
    Parameters
    ----------
    model:
        An ONNX model with parameters loaded in onnx package or path to file
        with models.
    mode: 'classifier', 'regressor' or None
        Mode of the converted coreml model:
        'classifier', a NeuralNetworkClassifier spec will be constructed.
        'regressor', a NeuralNetworkRegressor spec will be constructed.
    preprocessing_args:
        'is_bgr', 'red_bias', 'green_bias', 'blue_bias', 'gray_bias',
        'image_scale' keys with the same meaning as
        https://apple.github.io/coremltools/generated/coremltools.models.neural_network.html#coremltools.models.neural_network.NeuralNetworkBuilder.set_pre_processing_parameters
    deprocessing_args:
        Same as 'preprocessing_args' but for deprocessing.
    class_labels:
        As a string it represents the name of the file which contains
        the classification labels (one per line).
        As a list of strings it represents a list of categories that map
        the index of the output of a neural network to labels in a classifier.
    predicted_feature_name:
        Name of the output feature for the class labels exposed in the Core ML
        model (applies to classifiers only). Defaults to 'classLabel'
    add_custom_layers: bool
        Flag to turn on addition of custom CoreML layers for unsupported ONNX ops or attributes within
        a supported op.
    custom_conversion_functions: dict()
        A dictionary with keys corresponding to the names of onnx ops and values as functions taking
        an object of class 'Node' (see onnx-coreml/_graph.Node) and returning CoreML custom layer parameters.
    onnx_coreml_input_shape_map: dict()
        (Optional) A dictionary with keys corresponding to the model input names. Values are a list of integers that specify
        how the shape of the input is mapped to CoreML. Convention used for CoreML shapes is
        0: Sequence, 1: Batch, 2: channel, 3: height, 4: width.
        For example, an input of rank 2 could be mapped as [3,4] (i.e. H,W) or [1,2] (i.e. B,C) etc.
    Returns
    -------
    model: A coreml model.
    """
    if isinstance(model, Text):
        onnx_model = onnx.load(model)
    elif isinstance(model, onnx.ModelProto):
        onnx_model = model
    else:
        raise TypeError(
            "Model must be file path to .onnx file or onnx loaded model")

    transformers = [
        ConstantsToInitializers(),
        ShapeOpRemover(),
        ReshapeInitTensorFuser(),
        DropoutRemover(),
        UnsqueezeConstantRemover(),
        TransposeConstantRemover(),
        SliceConstantRemover(),
        ConcatConstantRemover(),
        ConvAddFuser(),
        BNBroadcastedMulFuser(),
        BNBroadcastedAddFuser(),
        PixelShuffleFuser(),
        AddModelInputsOutputs(),
        DivMulConstantRemover(),
        GatherConstantRemover(),
        ConstantFillToInitializers(),
    ]  # type: Iterable[Transformer]

    onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
    graph = _prepare_onnx_graph(onnx_model.graph, transformers)

    # are there ImageScaler nodes in the Graph?
    # If yes then add the info from it to the preprocessing dictionary, if the dictionary is not
    # already provided by the user
    if not bool(preprocessing_args):
        for node in graph.nodes:
            if node.op_type == 'ImageScaler':
                inp_name = node.inputs[0]
                scale = node.attrs.get('scale', 1.0)
                bias = node.attrs.get('bias', [0, 0, 0])
                if not (len(bias) == 1 or len(bias) == 3):
                    continue
                if 'image_scale' in preprocessing_args:
                    preprocessing_args['image_scale'][inp_name] = scale
                else:
                    preprocessing_args['image_scale'] = {inp_name: scale}
                if len(bias) == 3:
                    for i, color in enumerate(['red', 'green', 'blue']):
                        if color + '_bias' in preprocessing_args:
                            preprocessing_args[color +
                                               '_bias'][inp_name] = bias[i]
                        else:
                            preprocessing_args[color + '_bias'] = {
                                inp_name: bias[i]
                            }
                else:
                    if 'gray_bias' in preprocessing_args:
                        preprocessing_args['gray_bias'][inp_name] = bias[0]
                    else:
                        preprocessing_args['gray_bias'] = {inp_name: bias[0]}
                if inp_name not in image_input_names:
                    image_input_names.append(inp_name)  # type: ignore

    # remove all ImageScaler ops
    graph = graph.transformed([ImageScalerRemover()])

    #Make CoreML input and output features by gathering shape info and
    #interpreting it for CoreML
    input_features = _make_coreml_input_features(graph,
                                                 onnx_coreml_input_shape_map)
    if len(image_output_names) > 0:
        output_features = _make_coreml_output_features(graph, forceShape=True)
    else:
        output_features = _make_coreml_output_features(graph)

    builder = NeuralNetworkBuilder(input_features, output_features, mode=mode)
    _transform_coreml_dtypes(builder, graph.inputs, graph.outputs)

    is_deprocess_bgr_only = (len(deprocessing_args) == 1) and \
                            ("is_bgr" in deprocessing_args)
    add_deprocess = (len(image_output_names) > 0) and \
                    (len(deprocessing_args) > 0) and \
                    (not is_deprocess_bgr_only)

    if add_deprocess:
        mapping = {}
        for f in output_features:
            output_name = f[0]
            mapping[output_name] = graph.get_unique_edge_name(output_name)
        graph = OutputRenamer(mapping)(graph)

    if len(image_input_names) > 0:
        builder.set_pre_processing_parameters(
            image_input_names=image_input_names,
            is_bgr=preprocessing_args.get('is_bgr', False),
            red_bias=preprocessing_args.get('red_bias', 0.0),
            green_bias=preprocessing_args.get('green_bias', 0.0),
            blue_bias=preprocessing_args.get('blue_bias', 0.0),
            gray_bias=preprocessing_args.get('gray_bias', 0.0),
            image_scale=preprocessing_args.get('image_scale', 1.0))

    preprocessing_args.clear()

    if len(image_output_names) > 0:
        for f in output_features:
            f_name = f[0]
            if f_name in image_output_names:
                is_bgr = deprocessing_args.get('is_bgr', False)
                _convert_multiarray_output_to_image(builder.spec,
                                                    f_name,
                                                    is_bgr=is_bgr)
    '''Iterate through all the ops and translate them to CoreML layers.
    '''
    if not add_custom_layers:
        _check_unsupported_ops(graph.nodes)

    err = ErrorHandling(add_custom_layers, custom_conversion_functions)

    for i, node in enumerate(graph.nodes):
        print("%d/%d: Converting Node Type %s" %
              (i + 1, len(graph.nodes), node.op_type))
        _add_const_inputs_if_required(builder, node, graph, err)
        _convert_node(builder, node, graph, err)

    if DEBUG:
        plot_graph(graph,
                   graph_img_path='/tmp/after_conversion.pdf',
                   show_coreml_mapped_shapes=True)

    if add_deprocess:
        for f in output_features:
            output_name = f[0]
            if output_name not in image_output_names:
                continue
            output_shape = f[1].dimensions
            if len(output_shape) == 2 or output_shape[0] == 1:
                is_grayscale = True
            elif output_shape[0] == 3:
                is_grayscale = False
            else:
                raise ValueError('Output must be RGB image or Grayscale')
            _set_deprocessing(is_grayscale, builder, deprocessing_args,
                              mapping[output_name], output_name)

    if class_labels is not None:
        if isinstance(class_labels, Text):
            labels = [l.strip() for l in open(class_labels).readlines()
                      ]  # type: Sequence[Text]
        elif isinstance(class_labels, list):
            labels = class_labels
        else:
            raise TypeError("synset variable of unknown type. Type found: {}. \
                Expected either string or list of strings.".format(
                type(class_labels), ))

        builder.set_class_labels(class_labels=labels,
                                 predicted_feature_name=predicted_feature_name)

    # # add description to inputs/outputs that feed in/out of recurrent layers
    # for node_ in graph.nodes:
    #     if str(node_.op_type) in _SEQUENCE_LAYERS_REGISTRY:
    #         input_ = node_.inputs[0]
    #         output_ = node_.outputs[0]
    #         for i, inputs in enumerate(builder.spec.description.input):
    #             if inputs.name == input_:
    #                 builder.spec.description.input[i].shortDescription = 'This input is a sequence. '
    #         for i, outputs in enumerate(builder.spec.description.output):
    #             if outputs.name == output_:
    #                 builder.spec.description.output[i].shortDescription = 'This output is a sequence. '

    def _add_informative_description(feature, raise_error=True):
        if feature.type.WhichOneof('Type') == 'multiArrayType':
            if feature.name in graph.onnx_coreml_shape_mapping and feature.name in graph.shape_dict:
                mapp = graph.onnx_coreml_shape_mapping[feature.name]
                onnx_shape = graph.shape_dict[feature.name]
                if raise_error:
                    assert len(mapp) == len(
                        onnx_shape), "Something wrong in shape"
                if len(mapp) == len(onnx_shape):
                    shape = []
                    for i in range(5):
                        if i in mapp:
                            shape += [int(onnx_shape[mapp.index(i)])]
                        else:
                            shape += [1]
                    msg = 'MultiArray of shape {}. The first and second dimensions correspond to sequence and batch size, respectively'.format(
                        str(tuple(shape)))
                    feature.shortDescription += msg

    optional_input_names = []
    for tup in graph.optional_inputs:
        optional_input_names.append(tup[0])
    optional_output_names = []
    for tup in graph.optional_outputs:
        optional_output_names.append(tup[0])

    # add description for inputs and outputs shapes
    remove_input_id = []
    for i, input_ in enumerate(builder.spec.description.input):
        if input_.name not in optional_input_names:
            _add_informative_description(input_)
        else:
            remove_input_id.append(i)
    remove_output_id = []
    for i, output_ in enumerate(builder.spec.description.output):
        if output_.name not in optional_output_names:
            _add_informative_description(output_, raise_error=False)
        else:
            remove_output_id.append(i)

    for index in sorted(remove_input_id, reverse=True):
        del builder.spec.description.input[index]
    for index in sorted(remove_output_id, reverse=True):
        del builder.spec.description.output[index]

    if len(graph.optional_inputs) > 0 or len(graph.optional_outputs):
        builder.add_optionals(graph.optional_inputs, graph.optional_outputs)

    print(
        "Translation to CoreML spec completed. Now compiling the CoreML model."
    )
    try:
        if DEBUG:
            import coremltools
            coremltools.models.utils.save_spec(
                builder.spec, '/tmp/node_model_raw_spec.mlmodel')
        mlmodel = MLModel(builder.spec)
    except RuntimeError as e:
        raise ValueError('Compilation failed: {}'.format(str(e)))
    print('Model Compilation done.')

    # print information about all ops for which custom layers have been added
    if len(err.custom_layer_nodes) > 0:
        print('\n')
        print("Custom layers have been added to the CoreML model "
              "corresponding to the following ops in the onnx model: ")
        for i, node in enumerate(err.custom_layer_nodes):
            input_info = []
            for input_ in node.inputs:
                input_info.append(
                    (str(input_),
                     graph.shape_dict.get(input_, str("Shape not available"))))
            output_info = []
            for output_ in node.outputs:
                output_info.append(
                    (str(output_),
                     graph.shape_dict.get(output_,
                                          str("Shape not available"))))
            print(
                "{}/{}: op type: {}, op input names and shapes: {}, op output names and shapes: {}"
                .format(i + 1, len(err.custom_layer_nodes), node.op_type,
                        str(input_info), str(output_info)))

    return mlmodel
Exemplo n.º 10
0
def convert(
    model,  # type: Union[onnx.ModelProto, Text]
    mode=None,  # type: Optional[Text]
    image_input_names=[],  # type: Sequence[Text]
    preprocessing_args={},  # type: Dict[Text, Any]
    image_output_names=[],  # type: Sequence[Text]
    deprocessing_args={},  # type: Dict[Text, Any]
    class_labels=None,  # type: Union[Text, Iterable[Text], None]
    predicted_feature_name="classLabel",  # type: Text
    add_custom_layers=False,  # type: bool
    custom_conversion_functions={},  # type: Dict[Text, Any]
    onnx_coreml_input_shape_map={},  # type: Dict[Text, List[int,...]]
    minimum_ios_deployment_target="12",
):
    # type: (...) -> MLModel
    """
    WARNING: This function is deprecated. It will be removed in the 6.0.

    Convert ONNX model to CoreML.
    
    Parameters
    ----------
    model:
        An ONNX model with parameters loaded in the ONNX package, or path to file
        with models.
        
    mode: 'classifier', 'regressor' or None
    
        Mode of the converted coreml model:
        
        * ``'classifier'``: a NeuralNetworkClassifier spec will be constructed.
        * ``'regressor'``: a NeuralNetworkRegressor spec will be constructed.
        
    preprocessing_args:
        The ``'is_bgr'``, ``'red_bias'``, ``'green_bias'``, ``'blue_bias'``, ``'gray_bias'``,
        and ``'image_scale'`` keys have the same meaning as the pre-processing arguments for
        `NeuralNetworkBuilder <https://coremltools.readme.io/reference/modelsneural_network>`_.
    
    deprocessing_args:
        Same as ``'preprocessing_args'`` but for de-processing.
    
    class_labels:
        * As a string, it represents the name of the file which contains
          the classification labels (one per line).
        * As a list of strings, it represents a list of categories that map
          the index of the output of a neural network to labels in a classifier.
    
    predicted_feature_name:
        Name of the output feature for the class labels exposed in the Core ML
        model (applies to classifiers only). Defaults to ``'classLabel'``.
    
    add_custom_layers: bool
        Flag to turn on additional custom CoreML layers for unsupported ONNX ops or
    	attributes within a supported op.
    
    custom_conversion_functions: dict()
        * A dictionary with keys corresponding to the names/types of ONNX ops and values as 
          functions taking an object of the ``coreml-tools`` class:
          ``'NeuralNetworkBuilder'``, ``'Graph'`` (see ``onnx-coreml/_graph.Graph``),
          ``'Node'`` (see ``onnx-coreml/_graph.Node``), and 
          ``'ErrorHandling'`` (see ``onnx-coreml/_error_utils.ErrorHandling``).
        * This custom conversion function gets full control and responsibility for 
          converting a given ONNX op.
        * The function returns nothing and is responsible for adding an equivalent CoreML
          layer via ``'NeuralNetworkBuilder'``.
    
    onnx_coreml_input_shape_map: dict() (Optional) 
        * A dictionary with keys corresponding to the model input names.
        * Values are a list of integers that specify how the shape of the input is mapped
          to CoreML.
        * Convention used for CoreML shapes is ``0: Sequence``, ``1: Batch``,
          ``2: channel``, ``3: height``, ``4: width``. For example, an input of rank 2
          could be mapped as ``[3,4]`` (H,W) or ``[1,2]`` (B,C), and so on. This is
          ignored if ``minimum_ios_deployment_target`` is set to ``13``.
    
    minimum_ios_deployment_target: str
        Target Deployment iOS Version (default: ``'12'``). Supported iOS version options:
        ``'11.2'``, ``'12'``, ``'13'``. CoreML model produced by the converter will be
        compatible with the iOS version specified in this argument. For example, if
        ``minimum_ios_deployment_target = '12'``, the converter would utilize only CoreML
        features released up to version iOS12 (equivalent to macOS 10.14, watchOS 5, and
        so on). iOS 11.2 (CoreML 0.8) does not support ``resize_bilinear`` and
        ``crop_resize`` layers. See `supported v0.8 features <https://github.com/apple/coremltools/releases/tag/v0.8>`_.
        iOS 12 (CoreML 2.0), see `supported v2.0 features <https://github.com/apple/coremltools/releases/tag/v2.0>`_.
        iSO 13 (CoreML 3.0), see `supported v3.0 features <https://github.com/apple/coremltools/releases/tag/3.0-beta6>`_.
    
    
    Returns
    -------
    model: A coreml model.
    """
    if not _HAS_ONNX:
        raise ModuleNotFoundError("Missing ONNX package.")

    if isinstance(model, Text):
        onnx_model = onnx.load(model)
    elif isinstance(model, onnx.ModelProto):
        onnx_model = model
    else:
        raise TypeError(
            "Model must be file path to .onnx file or onnx loaded model")

    if not SupportedVersion.ios_support_check(minimum_ios_deployment_target):
        raise TypeError(
            "{} not supported. Please provide one of target iOS: {}",
            minimum_ios_deployment_target,
            SupportedVersion.get_supported_ios(),
        )

    global USE_SHAPE_MAPPING
    disable_coreml_rank5_mapping = False
    if SupportedVersion.is_nd_array_supported(minimum_ios_deployment_target):
        disable_coreml_rank5_mapping = True

    if disable_coreml_rank5_mapping:
        USE_SHAPE_MAPPING = False
    else:
        USE_SHAPE_MAPPING = True
    """
    First, apply a few optimizations to the ONNX graph,
    in preparation for conversion to CoreML.
    """

    # Using Dummy transformation to conditionally disable certain transformation
    class DummyTransformation(object):
        def __call__(self, graph):
            return graph

    transformers = [
        ConstantsToInitializers(),
        ShapeOpRemover(),
        ConstantRemover(),
        CastOpRemover(),
        PaddingOpRemover(),
        ReshapeInitTensorFuser(),
        DropoutRemover(),
        DeadCodeElimination(),
        ConvAddFuser(),
        BNBroadcastedMulFuser(),
        BNBroadcastedAddFuser(),
        ReshapeTransposeReshape_pattern1(),
        PixelShuffleFuser(),
        AddModelInputsOutputs()
        if not disable_coreml_rank5_mapping else DummyTransformation(),
        ConstantFillToInitializers(),
    ]  # type: Iterable[Transformer]

    onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
    graph = _prepare_onnx_graph(onnx_model.graph, transformers,
                                onnx_model.ir_version)
    """
    Check for ImageScalar nodes in ONNX, this will indicate whether input image preprocessing needs
    to be added to the CoreML graph or not.
    """
    # are there ImageScaler nodes in the Graph?
    # If yes then add the info from it to the "preprocessing_args" dictionary, if the dictionary is not
    # already provided by the user
    if not bool(preprocessing_args):
        for node in graph.nodes:
            if node.op_type == "ImageScaler":
                inp_name = node.inputs[0]
                scale = node.attrs.get("scale", 1.0)
                bias = node.attrs.get("bias", [0, 0, 0])
                if not (len(bias) == 1 or len(bias) == 3):
                    continue
                if "image_scale" in preprocessing_args:
                    preprocessing_args["image_scale"][inp_name] = scale
                else:
                    preprocessing_args["image_scale"] = {inp_name: scale}
                if len(bias) == 3:
                    for i, color in enumerate(["red", "green", "blue"]):
                        if color + "_bias" in preprocessing_args:
                            preprocessing_args[color +
                                               "_bias"][inp_name] = bias[i]
                        else:
                            preprocessing_args[color + "_bias"] = {
                                inp_name: bias[i]
                            }
                else:
                    if "gray_bias" in preprocessing_args:
                        preprocessing_args["gray_bias"][inp_name] = bias[0]
                    else:
                        preprocessing_args["gray_bias"] = {inp_name: bias[0]}
                if inp_name not in image_input_names:
                    image_input_names.append(inp_name)  # type: ignore

    # remove all ImageScaler ops
    graph = graph.transformed([ImageScalerRemover()])
    """
    Gather information (name, shape) for model inputs and outputs
    This information is then used to initialize the neural network builder object of coremltools.
    The builder object is later used to add layers to the CoreML model.
    """

    # Make CoreML input and output features by gathering shape info and
    # interpreting it for CoreML
    input_features = _make_coreml_input_features(graph,
                                                 onnx_coreml_input_shape_map,
                                                 disable_coreml_rank5_mapping)
    if len(image_output_names) > 0:
        output_features = _make_coreml_output_features(
            graph,
            forceShape=True,
            disable_coreml_rank5_mapping=disable_coreml_rank5_mapping,
        )
    else:
        output_features = _make_coreml_output_features(
            graph, disable_coreml_rank5_mapping=disable_coreml_rank5_mapping)

    builder = NeuralNetworkBuilder(
        input_features,
        output_features,
        mode=mode,
        disable_rank5_shape_mapping=disable_coreml_rank5_mapping,
    )

    # TODO: To be removed once, auto-downgrading of spec version is enabled
    builder.spec.specificationVersion = SupportedVersion.get_specification_version(
        minimum_ios_deployment_target)
    """
    Set CoreML input,output types (float, double, int) same as onnx types, if supported
    """
    _transform_coreml_dtypes(builder, graph.inputs, graph.outputs)
    """what follows is some book-keeping to support outputs of type image.
    """

    is_deprocess_bgr_only = (len(deprocessing_args)
                             == 1) and ("is_bgr" in deprocessing_args)
    add_deprocess = ((len(image_output_names) > 0)
                     and (len(deprocessing_args) > 0)
                     and (not is_deprocess_bgr_only))

    if add_deprocess:
        mapping = {}
        for f in output_features:
            output_name = f[0]
            mapping[output_name] = graph.get_unique_edge_name(output_name)
        graph = OutputRenamer(mapping)(graph)

    if len(image_input_names) > 0:
        builder.set_pre_processing_parameters(
            image_input_names=image_input_names,
            is_bgr=preprocessing_args.get("is_bgr", False),
            red_bias=preprocessing_args.get("red_bias", 0.0),
            green_bias=preprocessing_args.get("green_bias", 0.0),
            blue_bias=preprocessing_args.get("blue_bias", 0.0),
            gray_bias=preprocessing_args.get("gray_bias", 0.0),
            image_scale=preprocessing_args.get("image_scale", 1.0),
        )

    preprocessing_args.clear()

    if len(image_output_names) > 0:
        for f in output_features:
            f_name = f[0]
            if f_name in image_output_names:
                is_bgr = deprocessing_args.get("is_bgr", False)
                _convert_multiarray_output_to_image(builder.spec,
                                                    f_name,
                                                    is_bgr=is_bgr)
    """
    Iterate through all the ONNX ops and translate them to CoreML layers, one by one.
    """
    """
    before proceeding to start the layer translation process,
    check whether there is an op in the ONNX graph, whose translation function is not yet
    implemented in the converter or which is not supported in the CoreML framework. If so,
    raise an error before starting the process.
    (if the user desires to add a custom layer then this check is not required)
    """
    if not add_custom_layers:
        _check_unsupported_ops(graph.nodes, disable_coreml_rank5_mapping)
    """
    ErrorHandling is a generic class, useful to store a variety of parameters during the conversion process
    """
    err = ErrorHandling(add_custom_layers, custom_conversion_functions)

    for i, node in enumerate(graph.nodes):
        print("%d/%d: Converting Node Type %s" %
              (i + 1, len(graph.nodes), node.op_type))
        if disable_coreml_rank5_mapping:
            _convert_node_nd(builder, node, graph, err)
        else:
            _add_const_inputs_if_required(builder, node, graph, err)
            _convert_node(builder, node, graph, err)

    if DEBUG:
        plot_graph(
            graph,
            graph_img_path="/tmp/after_conversion.pdf",
            show_coreml_mapped_shapes=not disable_coreml_rank5_mapping,
        )

    if add_deprocess:
        for f in output_features:
            output_name = f[0]
            if output_name not in image_output_names:
                continue
            output_shape = f[1].dimensions
            if len(output_shape) == 2 or output_shape[0] == 1:
                is_grayscale = True
            elif output_shape[0] == 3:
                is_grayscale = False
            else:
                raise ValueError("Output must be RGB image or Grayscale")
            _set_deprocessing(
                is_grayscale,
                builder,
                deprocessing_args,
                mapping[output_name],
                output_name,
            )

    if class_labels is not None:
        if isinstance(class_labels, Text):
            labels = [l.strip() for l in open(class_labels).readlines()
                      ]  # type: Sequence[Text]
        elif isinstance(class_labels, list):
            labels = class_labels
        else:
            raise TypeError("synset variable of unknown type. Type found: {}. \
                Expected either string or list of strings.".format(
                type(class_labels), ))

        builder.set_class_labels(class_labels=labels,
                                 predicted_feature_name=predicted_feature_name)

    def _add_informative_description(feature, raise_error=True):
        if feature.type.WhichOneof("Type") == "multiArrayType":
            if (feature.name in graph.onnx_coreml_shape_mapping
                    and feature.name in graph.shape_dict):
                mapp = graph.onnx_coreml_shape_mapping[feature.name]
                onnx_shape = graph.shape_dict[feature.name]
                if raise_error:
                    assert len(mapp) == len(
                        onnx_shape), "Something wrong in shape"
                if len(mapp) == len(onnx_shape):
                    shape = []
                    for i in range(5):
                        if i in mapp:
                            shape += [int(onnx_shape[mapp.index(i)])]
                        else:
                            shape += [1]
                    msg = "MultiArray of shape {}. The first and second dimensions correspond to sequence and batch size, respectively".format(
                        str(tuple(shape)))
                    feature.shortDescription += msg

    optional_input_names = []
    for tup in graph.optional_inputs:
        optional_input_names.append(tup[0])
    optional_output_names = []
    for tup in graph.optional_outputs:
        optional_output_names.append(tup[0])

    # add description for inputs and outputs shapes
    remove_input_id = []
    for i, input_ in enumerate(builder.spec.description.input):
        if input_.name not in optional_input_names:
            if not disable_coreml_rank5_mapping:
                _add_informative_description(input_)
        else:
            remove_input_id.append(i)
    remove_output_id = []
    for i, output_ in enumerate(builder.spec.description.output):
        if output_.name not in optional_output_names:
            if not disable_coreml_rank5_mapping:
                _add_informative_description(output_, raise_error=False)
        else:
            remove_output_id.append(i)

    for index in sorted(remove_input_id, reverse=True):
        del builder.spec.description.input[index]
    for index in sorted(remove_output_id, reverse=True):
        del builder.spec.description.output[index]

    if len(graph.optional_inputs) > 0 or len(graph.optional_outputs):
        builder.add_optionals(graph.optional_inputs, graph.optional_outputs)

    # Check for specification version and target ios compatibility
    if (minimum_ios_deployment_target == "11.2"
            and builder.spec.WhichOneof("Type") == "neuralNetwork"):
        nn_spec = builder.spec.neuralNetwork
        for layer in nn_spec.layers:
            if (layer.WhichOneof("layer") == "resizeBilinear"
                    or layer.WhichOneof("layer") == "cropResize"):
                raise TypeError(
                    "{} not supported with target iOS 11.2 please provide higher target iOS"
                    .format(layer.WhichOneof("layer")))

    # Optimize ML Model Spec
    ml_model_passes = [remove_disconnected_layers, transform_conv_crop]
    for opt in ml_model_passes:
        opt(builder.spec)

    print(
        "Translation to CoreML spec completed. Now compiling the CoreML model."
    )
    try:
        if DEBUG:
            import coremltools

            coremltools.models.utils.save_spec(
                builder.spec, "/tmp/node_model_raw_spec.mlmodel")
            from coremltools.models.neural_network.printer import print_network_spec

            print_network_spec(builder.spec, style="coding")
        mlmodel = MLModel(builder.spec)
    except RuntimeError as e:
        raise ValueError("Compilation failed: {}".format(str(e)))
    print("Model Compilation done.")

    # print information about all ops for which custom layers have been added
    if len(err.custom_layer_nodes) > 0:
        print("\n")
        print("Custom layers have been added to the CoreML model "
              "corresponding to the following ops in the onnx model: ")
        for i, node in enumerate(err.custom_layer_nodes):
            input_info = []
            for input_ in node.inputs:
                input_info.append((
                    str(input_),
                    graph.shape_dict.get(input_, str("Shape not available")),
                ))
            output_info = []
            for output_ in node.outputs:
                output_info.append((
                    str(output_),
                    graph.shape_dict.get(output_, str("Shape not available")),
                ))
            print(
                "{}/{}: op type: {}, op input names and shapes: {}, op output names and shapes: {}"
                .format(
                    i + 1,
                    len(err.custom_layer_nodes),
                    node.op_type,
                    str(input_info),
                    str(output_info),
                ))

    mlmodel.user_defined_metadata[_METADATA_VERSION] = ct_version
    mlmodel.user_defined_metadata[_METADATA_SOURCE] = "onnx=={0}".format(
        onnx.__version__)
    return mlmodel
Exemplo n.º 11
0
def convert(model,
            input_shapes,
            input_names=['input'],
            output_names=['output'],
            mode=None,
            image_input_names=[],
            preprocessing_args={},
            image_output_names=[],
            deprocessing_args={},
            class_labels=None,
            predicted_feature_name='classLabel',
            unknown_layer_converter_fn=None):
    """
    Convert Torch7 model to CoreML.

    Parameters
    ----------
    model: Torch7 model (loaded with PyTorch) | str
        A trained Torch7 model loaded in python using PyTorch or path to file
        with model (*.t7).

    input_shapes: list of tuples
        Shapes of the input tensors.

    mode: str ('classifier', 'regressor' or None)
        Mode of the converted coreml model:
        'classifier', a NeuralNetworkClassifier spec will be constructed.
        'regressor', a NeuralNetworkRegressor spec will be constructed.

    preprocessing_args: dict
        'is_bgr', 'red_bias', 'green_bias', 'blue_bias', 'gray_bias',
        'image_scale' keys with the same meaning as
        https://apple.github.io/coremltools/generated/coremltools.models.neural_network.html#coremltools.models.neural_network.NeuralNetworkBuilder.set_pre_processing_parameters

    deprocessing_args: dict
        Same as 'preprocessing_args' but for deprocessing.

    class_labels: A string or list of strings.
        As a string it represents the name of the file which contains
        the classification labels (one per line).
        As a list of strings it represents a list of categories that map
        the index of the output of a neural network to labels in a classifier.

    predicted_feature_name: str
        Name of the output feature for the class labels exposed in the Core ML
        model (applies to classifiers only). Defaults to 'classLabel'

    unknown_layer_converter_fn: function with signature:
        (builder, name, layer, input_names, output_names)
            builder: object - instance of NeuralNetworkBuilder class
            name: str - generated layer name
            layer: object - pytorch object for corresponding layer
            input_names: list of strings
            output_names: list of strings
            Returns: list of strings for layer output names
        Callback function to handle unknown for torch2coreml layers


    Returns
    -------
    model: A coreml model.
    """
    _gen_layer_name.called = 0
    _get_layer_converter_fn.unknown_converter_fn = unknown_layer_converter_fn

    if isinstance(model, basestring):
        torch_model = load_lua(model)
    elif isinstance(model, torch.legacy.nn.Sequential):
        torch_model = model
    else:
        raise TypeError(
            "Model must be file path to .t7 file or pytorch loaded model \
            with torch.legacy.nn.Sequential module as root")

    torch_model.evaluate()

    if not isinstance(input_shapes, list):
        raise TypeError("Input shapes should be a list of tuples.")

    for shape in input_shapes:
        if not isinstance(shape, tuple):
            raise TypeError("Input shape should be a tuple.")

    if len(input_names) != len(input_shapes):
        raise ValueError(
            "Input names count must be equal to input shapes count")

    output_shapes = _infer_torch_output_shapes(torch_model, input_shapes)

    if len(output_shapes) != len(output_names):
        raise ValueError(
            "Model has {} outputs, but you set output_names for {}.".format(
                len(output_shapes), len(output_names)))

    # create input/output features
    input_features = []
    for i in range(len(input_names)):
        input_features.append(
            (input_names[i], datatypes.Array(*input_shapes[i])))
    output_features = []
    for i in range(len(output_names)):
        output_features.append(
            (output_names[i], datatypes.Array(*output_shapes[i])))

    builder = NeuralNetworkBuilder(input_features, output_features, mode)

    # build model
    layer_name = _gen_layer_name(torch_model)
    _output_names = output_names[:]
    if len(image_output_names) > 0:
        for i in range(len(_output_names)):
            if _output_names[i] in image_output_names:
                _output_names[i] = _gen_layer_name(_DEPROCESS_LAYER_NAME)

    model_output_names = _layers._convert_layer(builder, layer_name,
                                                torch_model, input_names,
                                                _output_names)

    # set preprocessing parameters
    if len(image_input_names) > 0:
        builder.set_pre_processing_parameters(
            image_input_names=image_input_names,
            is_bgr=preprocessing_args.get('is_bgr', False),
            red_bias=preprocessing_args.get('red_bias', 0.0),
            green_bias=preprocessing_args.get('green_bias', 0.0),
            blue_bias=preprocessing_args.get('blue_bias', 0.0),
            gray_bias=preprocessing_args.get('gray_bias', 0.0),
            image_scale=preprocessing_args.get('image_scale', 1.0))

    # set deprocessing parameters
    if len(image_output_names) > 0:
        for i in range(len(output_names)):
            output_name = output_names[i]
            if output_name in image_output_names:
                output_shape = output_shapes[i]
                if len(output_shape) == 2 or output_shape[0] == 1:
                    is_grayscale = True
                elif output_shape[0] == 3:
                    is_grayscale = False
                else:
                    raise ValueError('Output must be RGB image or Grayscale')
                _set_deprocessing(is_grayscale, builder, deprocessing_args,
                                  model_output_names[i], output_name)

    if class_labels is not None:
        if type(class_labels) is str:
            labels = [l.strip() for l in open(class_labels).readlines()]
        elif type(class_labels) is list:
            labels = class_labels
        else:
            raise TypeError("synset variable of unknown type. Type found: {}. \
                Expected either string or list of strings.".format(
                type(class_labels), ))

        builder.set_class_labels(class_labels=labels,
                                 predicted_feature_name=predicted_feature_name)

    return MLModel(builder.spec)