def _add_mean(op, context): input_name = compat.as_bytes(op.inputs[0].name) output_name = compat.as_bytes(op.outputs[0].name) axis_ind = context.consts[op.inputs[1].name] input_shape = context.shape_dict[input_name] if context.use_dfs_shape_infer: status = interpret_shape(input_name, context) else: status = False if status: labeled_shape = context.dim_labels[input_name] if isinstance(axis_ind, np.ndarray): axis = '' for i in axis_ind: if input_shape[i] != 1: axis += labeled_shape[i] axis = ''.join(sorted(axis)) else: axis = labeled_shape[axis_ind] assert axis in [ 'S', 'C', 'H', 'W', 'CHW', 'HW' ], ('Axis value %s not supported. ' 'Reduction supported along C, H, W, HW, CHW dimensions only.' % axis) else: if len(input_shape) == 4 and ( np.array_equal(axis_ind, np.array([0, 1, 2])) or np.array_equal(axis_ind, np.array([1, 2]))): axis = 'HW' else: assert False, 'Mean axis case not handled currently' mode = 'avg' # The simple case; reduction along non sequence axis if axis != 'S': context.builder.add_reduce(output_name, input_name, output_name, axis, mode) # Need to permute, reduce and then permute back else: context.builder.add_permute(output_name, (1, 0, 2, 3), input_name, output_name + '_swap_Seq_C') context.builder.add_reduce(output_name, output_name + '_swap_Seq_C', output_name + '_pre_permute', 'C', mode) context.builder.add_permute(output_name, (1, 0, 2, 3), output_name + '_pre_permute', output_name) context.translated[output_name] = True
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)
def _add_concat(op, context): output_name = compat.as_bytes(op.outputs[0].name) output_shape = context.shape_dict[output_name] axis = 3 #3 -> 'Channel', 2 -> 'Width', 1 -> 'Height if op.type == 'Concat': axis_name = compat.as_bytes(op.inputs[0].name) axis = context.consts[axis_name] input_names = [] for i, input in enumerate(op.inputs): if i == 0: continue input_names.append(compat.as_bytes(input.name)) if op.type == 'ConcatV2': axis_name = compat.as_bytes(op.inputs[-1].name) axis = context.consts[axis_name] input_names = [] for i, input in enumerate(op.inputs): if i == len(op.inputs) - 1: continue input_names.append(compat.as_bytes(input.name)) if context.use_dfs_shape_infer: status = interpret_shape(output_name, context) else: status = False if status: labeled_shape = context.dim_labels[output_name] if labeled_shape[axis] == 'C': axis = 3 elif labeled_shape[axis] == 'H': axis = 1 elif labeled_shape[axis] == 'W': axis = 2 else: assert False, 'Concatenation supported only along channel, height or '\ 'width dimensions' else: if len(output_shape) == 4: assert axis in [1, 2, 3], 'Concat axis case not handled' elif len(output_shape) == 3: axis += 1 elif len(output_shape) == 1: axis = 3 else: assert False, 'Concat axis case not handled' # Temporary workaround for fixing bugs on certain devices. # TODO: remove this in future # If concat's input is coming from another pool/concat: insert a linear activation layer, # if it hasn't been inserted already coreml_layers = context.builder.nn_spec.layers coreml_outputs = dict() for layer in coreml_layers: for out in layer.output: coreml_outputs[out] = True for layer in coreml_layers: if layer.WhichOneof('layer') in ['concat', 'pooling']: for i, inp in enumerate(input_names): if layer.output[0] == inp: out = inp + '__linear_activation' if out not in coreml_outputs: context.builder.add_activation(out, 'LINEAR', inp, out, [1.0, 0]) input_names[i] = out if axis == 3: #concatenate along channel axis context.builder.add_elementwise(output_name, input_names, output_name, 'CONCAT') elif axis == 2: #concatenate along width axis blob_postfix = '_swap_W_C_' transpose_order = (0, 3, 2, 1) inputs_permuted = [] for i, input_name in enumerate(input_names): context.builder.add_permute(output_name + '_' + str(i), transpose_order, input_name, input_name + blob_postfix + str(i)) inputs_permuted.append(input_name + blob_postfix + str(i)) context.builder.add_elementwise(output_name + '_concat', inputs_permuted, output_name + '_concat', 'CONCAT') context.builder.add_permute(output_name, transpose_order, output_name + '_concat', output_name) elif axis == 1: #concatenate along height axis inputs_permuted = [] for i, input_name in enumerate(input_names): context.builder.add_permute(output_name + '_' + str(i), (0, 2, 1, 3), input_name, input_name + '_swap_H_C_' + str(i)) inputs_permuted.append(input_name + '_swap_H_C_' + str(i)) context.builder.add_elementwise(output_name + '_concat', inputs_permuted, output_name + '_concat', 'CONCAT') context.builder.add_permute(output_name, (0, 2, 1, 3), output_name + '_concat', output_name) else: assert False, 'Concat axis case not handled' context.translated[output_name] = True
def _add_reduce(op, context, mode): input_name = compat.as_bytes(op.inputs[0].name) output_name = compat.as_bytes(op.outputs[0].name) axis_ind = context.consts[op.inputs[1].name] input_shape = context.shape_dict[input_name] if context.use_dfs_shape_infer: status = interpret_shape(input_name, context) else: status = False # Determine reduction axis labels axis = None if status: labeled_shape = context.dim_labels[input_name] if isinstance(axis_ind, np.ndarray): axis = '' for i in axis_ind: if input_shape[i] != 1: axis += labeled_shape[i] axis = ''.join(sorted(axis)) else: axis = labeled_shape[axis_ind] assert axis in [ 'S', 'C', 'H', 'W', 'CHW', 'HW' ], ('Axis value %s not supported. ' 'Reduction supported along C, H, W, HW, CHW dimensions only.' % axis) else: if isinstance(axis_ind, np.ndarray): axis_ind = axis_ind.tolist() if len(axis_ind) == 1: axis_ind = axis_ind[0] elif len(input_shape) == len(axis_ind): axis = 'CHW' if axis is None: # single axis reduction axis_ind = (len(input_shape) + axis_ind) if axis_ind < 0 else axis_ind if len(input_shape) == 4: if axis_ind == 3: axis = 'C' elif len(input_shape) == 2: if axis_ind == 0: # TODO - only works for stylenet. (W,C)--->(1,C) axis = 'W' elif axis_ind == 1: axis = 'C' elif len(input_shape) == 1: if axis_ind == 0: axis = 'CHW' elif len(input_shape) == 3: if axis_ind == 2: axis = 'C' if axis == None: raise NotImplementedError( 'Reduce axis %s for input shape %s not handled currently' % (str(axis_ind), str(input_shape))) # The simple case; reduction along non sequence axis if axis != 'S': context.builder.add_reduce(output_name, input_name, output_name, axis, mode) # Need to permute, reduce and then permute back else: context.builder.add_permute(output_name, (1, 0, 2, 3), input_name, output_name + '_swap_Seq_C') context.builder.add_reduce(output_name, output_name + '_swap_Seq_C', output_name + '_pre_permute', 'C', mode) context.builder.add_permute(output_name, (1, 0, 2, 3), output_name + '_pre_permute', output_name) context.translated[output_name] = True
def _add_reshape(op, context): input_name = compat.as_bytes(op.inputs[0].name) output_name = compat.as_bytes(op.outputs[0].name) #First make sure the the input blob exists in the CoreML graph input_name = _layers.make_tensor(op.inputs[0], context) input_shape = context.shape_dict[input_name] target_shape = context.shape_dict[output_name] squeezed_input_shape = _remove_beginning_unit_dimensions(input_shape) squeezed_output_shape = _remove_beginning_unit_dimensions(target_shape) if squeezed_input_shape == squeezed_output_shape: # reshape is either squeeze or expand_dim _layers.skip(op, context) return if context.use_dfs_shape_infer: status = interpret_shape(output_name, context) else: status = False if status: target_shape = context.shape_dict_rank_4[output_name] if interpret_shape(input_name, context): input_shape_rank_4 = context.shape_dict_rank_4[input_name] if input_shape_rank_4 == target_shape: _layers.skip(op, context) return # When reshape is immediately followed by squeeze if len(op.outputs) > 0 and len(op.outputs[0].consumers()) > 0 and \ op.outputs[0].consumers()[0].type == 'Squeeze': squeezed_output_name = compat.as_bytes( op.outputs[0].consumers()[0].outputs[0].name) target_shape = context.shape_dict[squeezed_output_name] # TODO - these cases of reshape are just for mobilenet and stylenet: # if target_shape == (1,X) ----> new_shape = (X,1,1) # if target_shape == (X,1) -----> new_shape = (1,1,X) assert len(target_shape) in [ 1, 2, 3, 4 ], ('Reshape: Currently only supported if target shape is rank 2, 3 or 4') mode = 0 if len(target_shape) == 2: if target_shape[1] != 1: #(1,X) new_shape = (1, target_shape[1], 1, 1) if len(input_shape) == 4 or len(input_shape) == 3: # (N,H,W,C) --> (1,C) or (N,S,C) --> (N,1,W,C) mode = 1 else: new_shape = (1, 1, 1, target_shape[0]) context.builder.add_reshape(output_name, input_name, output_name, new_shape, mode) elif len(target_shape) == 3: # Target shape is [H,W,C] --> [1, C, H, W] new_shape = (1, target_shape[2], target_shape[0], target_shape[1]) context.builder.add_reshape(output_name, input_name, output_name, new_shape, 1) elif len(target_shape) == 4: new_shape = (target_shape[0], target_shape[3], target_shape[1], target_shape[2]) context.builder.add_reshape(output_name, input_name, output_name, new_shape, 1) elif len(target_shape) == 1: new_shape = (1, target_shape[0], 1, 1) context.builder.add_reshape(output_name, input_name, output_name, new_shape, 1) context.translated[output_name] = True
def _add_const(context, name, x, output_name, shape=None): if output_name in context.load_constants_mlmodel: return if shape is not None: context.builder.add_load_constant(name, output_name, x, shape) context.load_constants_mlmodel[output_name] = True return context.load_constants_mlmodel[output_name] = True if context.use_dfs_shape_infer: status = interpret_shape(output_name, context) else: status = False if status: rank_4_shape = context.shape_dict_rank_4[output_name] # TODO - Interpreting 1st dimension as seq. in this case instead of batch seq, h, w, c = rank_4_shape x = np.reshape(x, (seq, h, w, c)) #first check the simple case: seq. dimension is 1 if seq == 1: shape = [c, h, w] # (C, H, W) x = np.transpose(x, [0, 3, 1, 2]) context.builder.add_load_constant(name, output_name, x, shape) #when sequence dimension is not 1, we need some permute layers as well #since CoreML only allows loading constant of rank-3: [C,H,W]) else: assert c == 1 or h == 1 or w == 1, \ 'Add constant: cannot add a constant in which all the dimensions ' \ '(Seq, C, H, W) are of non-unit size' if c == 1: #swap seq. and C x = np.transpose(x, [3, 0, 1, 2]) #(S,H,W,C) --> (C,S,H,W) context.builder.add_load_constant(name + '_pre_permute', output_name + '_pre_permute', x, [seq, h, w]) context.builder.add_permute(output_name, (1, 0, 2, 3), output_name + '_pre_permute', output_name) elif h == 1: #swap seq. and H x = np.transpose(x, [1, 3, 0, 2]) #(S,H,W,C) --> (H,C,S,W) context.builder.add_load_constant(name + '_pre_permute', output_name + '_pre_permute', x, [c, seq, w]) context.builder.add_permute(output_name, (2, 1, 0, 3), output_name + '_pre_permute', output_name) else: # w == 1, swap seq. and W x = np.transpose(x, [2, 3, 1, 0]) #(S,H,W,C) --> (W,C,H,S) context.builder.add_load_constant(name + '_pre_permute', output_name + '_pre_permute', x, [c, h, seq]) context.builder.add_permute(output_name, (3, 1, 2, 0), output_name + '_pre_permute', output_name) else: #Static shape mapping shape = list(x.shape) assert len(shape) < 5, 'Const blob shape is more than rank 4' if len(shape) == 0: shape = [1, 1, 1] #(1,1,1) elif len(shape) == 1: shape = [shape[0], 1, 1] #(C,1,1) elif len(shape) == 2: shape = [ shape[1], 1, shape[0] ] # HACK: (W,C) ---> (C,1,W) . Style transfer matrices are (W,C) x = np.transpose(x, [1, 0]) elif len(shape) == 3: shape = [shape[2], shape[0], shape[1]] # (H,W,C) ---> (C,H,W) x = np.transpose(x, [2, 0, 1]) elif len(shape) == 4: assert shape[0] == 1, 'Add Constant: Batch dimension must be 1' shape = [shape[3], shape[1], shape[2]] #(B,H,W,C) ---> (C,H,W) x = x[0, :, :, :] #(H,W,C) x = np.transpose(x, [2, 0, 1]) context.builder.add_load_constant(name, output_name, x, shape)