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)
