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
def train_model(ENV, in_file, op_file): graph = tf.Graph() with graph.as_default(): stacked_layers = {} # e.g: log filter bank or MFCC features # Has size [batch_size, max_stepsize, num_features], but the # batch_size and max_stepsize can vary along each step inputs = tf.placeholder(tf.float32, [None, None, num_features]) targets = tf.sparse_placeholder(tf.int32) # 1d array of size [batch_size] seq_len = tf.placeholder(tf.int32, [None]) # Weights & biases weight_classes = tf.Variable( tf.truncated_normal([num_hidden, num_classes], mean=0, stddev=0.1, dtype=tf.float32)) bias_classes = tf.Variable(tf.zeros([num_classes]), dtype=tf.float32) #_activation = tf.nn.relu#this was causing the model to diverge _activation = None layers = {'forward': [], 'backward': []} for key in layers.keys(): for i in range(num_layers): cell = tf.nn.rnn_cell.LSTMCell(num_hidden, use_peepholes=True, activation=_activation, state_is_tuple=True, cell_clip=clip_thresh) # #cell = RWACell(num_units=num_hidden) layers[key].append(cell) stacked_layers[key] = tf.nn.rnn_cell.MultiRNNCell( layers[key], state_is_tuple=True) outputs, bilstm_vars = tf.nn.bidirectional_dynamic_rnn( stacked_layers['forward'], stacked_layers['backward'], inputs, sequence_length=seq_len, time_major=False, # [batch_size, max_time, num_hidden] dtype=tf.float32) """ outputs_concate = tf.concat_v2(outputs, 2) outputs_concate = tf.reshape(outputs_concate, [-1, 2*num_hidden]) # logits = tf.matmul(outputs_concate, weight_classes) + bias_classes """ fw_output = tf.reshape(outputs[0], [-1, num_hidden]) bw_output = tf.reshape(outputs[1], [-1, num_hidden]) logits = tf.add( tf.add(tf.matmul(fw_output, weight_classes), tf.matmul(bw_output, weight_classes)), bias_classes) logits = tf.reshape(logits, [batch_size, -1, num_classes]) loss = tf.nn.ctc_loss(targets, logits, seq_len, time_major=False) error = tf.reduce_mean(loss) optimizer = tf.train.MomentumOptimizer(learning_rate, momentum).minimize(error) # Evaluating # decoded, log_prob = ctc_ops.ctc_greedy_decoder(tf.transpose(logits, perm=[1, 0, 2]), seq_len) decoded, log_prob = tf.nn.ctc_beam_search_decoder( tf.transpose(logits, perm=[1, 0, 2]), seq_len) label_error_rate = tf.reduce_mean( tf.edit_distance(tf.cast(decoded[0], tf.int32), targets)) gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5) data, labels = load_ipad_data(in_file) bound = ((3 * len(data) / batch_size) / 4) * batch_size train_inputs = data[0:bound] train_labels = labels[0:bound] test_data = data[bound:] test_labels = labels[bound:] num_examples = len(train_inputs) num_batches_per_epoch = num_examples / batch_size with tf.Session(graph=graph, config=tf.ConfigProto(gpu_options=gpu_options)) as session: # Initializate the weights and biases tf.global_variables_initializer().run() saver = tf.train.Saver(tf.global_variables(), max_to_keep=0) ckpt = tf.train.get_checkpoint_state(op_file) if ckpt: logging.info('load', ckpt.model_checkpoint_path) saver.restore(session, ckpt.model_checkpoint_path) else: logging.info("no previous session to load") for curr_epoch in range(num_epochs): train_cost = train_ler = 0 start = time.time() for batch in range(num_batches_per_epoch): # Getting the index indices = [ i % num_examples for i in range(batch * batch_size, (batch + 1) * batch_size) ] batch_train_inputs = train_inputs[indices] # Padding input to max_time_step of this batch batch_train_inputs, batch_train_seq_len = pad_sequences( batch_train_inputs) # Converting to sparse representation so as to to feed SparseTensor input batch_train_targets = sparse_tuple_from(train_labels[indices]) feed = { inputs: batch_train_inputs, targets: batch_train_targets, seq_len: batch_train_seq_len } batch_cost, _ = session.run([error, optimizer], feed) train_cost += batch_cost * batch_size train_ler += session.run(label_error_rate, feed_dict=feed) * batch_size log = "Epoch {}/{}, iter {}, batch_cost {}" logging.info( log.format(curr_epoch + 1, num_epochs, batch, batch_cost)) saver.save(session, os.path.join(ENV.output, 'best.ckpt'), global_step=curr_epoch) # Shuffle the data shuffled_indexes = np.random.permutation(num_examples) train_inputs = train_inputs[shuffled_indexes] train_labels = train_labels[shuffled_indexes] # Metrics mean train_cost /= num_examples train_ler /= num_examples log = "Epoch {}/{}, train_cost = {:.3f}, train_ler = {:.3f}, time = {:.3f}" logging.info( log.format(curr_epoch + 1, num_epochs, train_cost, train_ler, time.time() - start)) #run the test data through indices = [ i % len(test_data) for i in range(batch * batch_size, (batch + 1) * batch_size) ] test_inputs = test_data[indices] test_inputs, test_seq_len = pad_sequences(test_inputs) test_targets = sparse_tuple_from(test_labels[indices]) feed_test = { inputs: test_inputs, targets: test_targets, seq_len: test_seq_len } test_cost, test_ler = session.run([error, label_error_rate], feed_dict=feed_test) log = "Epoch {}/{}, test_cost {}, test_ler {}" logging.info( log.format(curr_epoch + 1, num_epochs, test_cost, test_ler)) input_features = [('strokeData', datatypes.Array(num_features))] output_features = [('labels', datatypes.Array(num_classes))] vars = tf.trainable_variables() weights = {'forward': {}, 'backward': {}} for _var in vars: name = _var.name.encode('utf-8') if name.startswith('bidirectional_rnn/fw'): key = name.replace('bidirectional_rnn/fw/', '') key = key.replace('multi_rnn_cell/cell_0/lstm_cell/', '') key = key.replace(':0', '') weights['forward'][key] = _var.eval() else: key = name.replace('bidirectional_rnn/bw/', '') key = key.replace('multi_rnn_cell/cell_0/lstm_cell/', '') key = key.replace(':0', '') weights['backward'][key] = _var.eval() builder = NeuralNetworkBuilder(input_features, output_features, mode=None) fw_biases = [ weights['forward']['bias'][0 * num_hidden:1 * num_hidden], weights['forward']['bias'][1 * num_hidden:2 * num_hidden], weights['forward']['bias'][2 * num_hidden:3 * num_hidden], weights['forward']['bias'][3 * num_hidden:4 * num_hidden] ] bw_biases = [ weights['backward']['bias'][0 * num_hidden:1 * num_hidden], weights['backward']['bias'][1 * num_hidden:2 * num_hidden], weights['backward']['bias'][2 * num_hidden:3 * num_hidden], weights['backward']['bias'][3 * num_hidden:4 * num_hidden] ] num_LSTM_gates = 5 input_weights = { 'forward': np.zeros((num_LSTM_gates - 1, num_hidden, num_features)), 'backward': np.zeros((num_LSTM_gates - 1, num_hidden, num_features)) } recurrent_weights = { 'forward': np.zeros((num_LSTM_gates - 1, num_hidden, num_hidden)), 'backward': np.zeros((num_LSTM_gates - 1, num_hidden, num_hidden)) } builder.add_bidirlstm( name='bidirectional_1', W_h=recurrent_weights['forward'], W_x=input_weights['forward'], b=fw_biases, W_h_back=recurrent_weights['backward'], W_x_back=input_weights['backward'], b_back=bw_biases, hidden_size=num_hidden, input_size=num_features, input_names=[ 'strokeData', 'bidirectional_1_h_in', 'bidirectional_1_c_in', 'bidirectional_1_h_in_rev', 'bidirectional_1_c_in_rev' ], output_names=[ 'y', 'bidirectional_1_h_out', 'bidirectional_1_c_out', 'bidirectional_1_h_out_rev', 'bidirectional_1_c_out_rev' ], peep=[ weights['forward']['w_i_diag'], weights['forward']['w_f_diag'], weights['forward']['w_o_diag'] ], peep_back=[ weights['backward']['w_i_diag'], weights['backward']['w_f_diag'], weights['backward']['w_o_diag'] ], cell_clip_threshold=clip_thresh) builder.add_softmax(name='softmax', input_name='y', output_name='labels') optional_inputs = [('bidirectional_1_h_in', num_hidden), ('bidirectional_1_c_in', num_hidden), ('bidirectional_1_h_in_rev', num_hidden), ('bidirectional_1_c_in_rev', num_hidden)] optional_outputs = [('bidirectional_1_h_out', num_hidden), ('bidirectional_1_c_out', num_hidden), ('bidirectional_1_h_out_rev', num_hidden), ('bidirectional_1_c_out_rev', num_hidden)] #not really sure what this line belowe does, just copied it from the Keras converter in coremltools, # and it seemed to make things work builder.add_optionals(optional_inputs, optional_outputs) model = MLModel(builder.spec) model.short_description = 'Model for recognizing a symbols and diagrams drawn on ipad screen with apple pencil' model.input_description[ 'strokeData'] = 'A collection of strokes to classify' model.output_description[ 'labels'] = 'The "probability" of each label, in a dense array' outfile = 'bilstm.mlmodel' model.save(outfile) print('Saved to file: %s' % outfile)
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
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