<jupyter_start><jupyter_text>This notebook demonstrates the process of adding custom layers to the CoreML model during conversion. We discuss  three examples.For TensorFlow operations (ops for short) that are not translatable to any of the CoreML layers, custom layers can be inserted in the CoreML model (list of CoreML layers can be found [here](https://github.com/apple/coremltools/blob/master/mlmodel/format/NeuralNetwork.proto) or [here](https://apple.github.io/coremltools/coremlspecification/sections/NeuralNetwork.html)). At runtime, CoreML framework will look for the implementation code of the custom layers, which has to be provided by the developer in their app.   Custom layer is a [proto message](https://github.com/apple/coremltools/blob/5b5b8190764ffe78110be6b4d0edbeebe0253a6e/mlmodel/format/NeuralNetwork.protoL2280), like any other neural network layer in the .mlmodel file (which is in the protobuf format), that can hold the parameters and weights (if any) associated with the TF op.Here is the [documentation](https://developer.apple.com/documentation/coreml/core_ml_api/creating_a_custom_layer) on CoreML custom layers and a nice detailed [blogpost](http://machinethink.net/blog/coreml-custom-layers/). There are two ways in which a custom layer can be added during conversion from TF:1. Specify the argument "add_custom_layers=True" during conversion. This will automatically check for unsupported ops and insert a coreml custom layer message in place of that op. The message can be later edited, if required, to add/remove any parameters.            2. Specify the arguments "add_custom_layers=True" and "custom_conversion_functions" to the converter. The second argument is a dictionary, with keys that are either op types or op names and values are user-defined function handles. The functions receive TensorFlow [op](https://github.com/tensorflow/tensorflow/blob/51ef16057b4625e0a3e2943a9f1bbf856cf098ca/tensorflow/python/framework/ops.pyL3707) object and the CoreML neural network [builder object](https://github.com/apple/coremltools/blob/5b5b8190764ffe78110be6b4d0edbeebe0253a6e/coremltools/models/neural_network.pyL34) and give the user full control on how to handle the TF op and which layers to add to the CoreML graph. When the key is an op type, the function is called whenever op of that type is encountered while traversing the TF graph. Operation names as keys are useful for targeting specific ops. Lets now dive into the examples to make this process clear. First up, setting up some utilities:<jupyter_code>from __future__ import print_function
import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.python.tools.freeze_graph import freeze_graph
import numpy as np
import shutil
import tempfile
import os
import tfcoreml
import coremltools
from coremltools.proto import NeuralNetwork_pb2
import netron # we use netron: https://github.com/lutzroeder/Netron for visualization. Comment out this line and all the calls to the "_visualize" method, if you do not want to use it. 
# A utility function to freeze rhe graph. It will be used later
def _simple_run_and_freeze(graph, output_name, frozen_model_file='', feed_dict={}):
    
    model_dir = tempfile.mkdtemp()
    graph_def_file = os.path.join(model_dir, 'tf_graph.pbtxt')
    checkpoint_file = os.path.join(model_dir, 'tf_model.ckpt')
    
    tf.reset_default_graph()
    with graph.as_default() as g:
      saver = tf.train.Saver()

    with tf.Session(graph = graph) as sess:
      # initialize
      sess.run(tf.global_variables_initializer())
      # run the result
      fetch = graph.get_operation_by_name(output_name).outputs[0]
      tf_result = sess.run(fetch, feed_dict=feed_dict)
      # save graph definition somewhere
      tf.train.write_graph(sess.graph, model_dir, graph_def_file)
      # save the weights
      saver.save(sess, checkpoint_file)
    
    freeze_graph(input_graph=graph_def_file,
                 input_saver="",
                 input_binary=False,
                 input_checkpoint=checkpoint_file,
                 output_node_names=output_name,
                 restore_op_name="save/restore_all",
                 filename_tensor_name="save/Const:0",
                 output_graph=frozen_model_file,
                 clear_devices=True,
                 initializer_nodes="")
    
    if os.path.exists(model_dir):
        shutil.rmtree(model_dir)
    
    return tf_result

# A utility function that takes an MLModel instance and prints info about Neural network layers inside.
# It prints short info about all the NN layers and the full description of any custom layer found
def _print_coreml_nn_layer_info(spec):
    nn_layers = coremltools.models.utils._get_nn_layers(spec)
    for i, layer in enumerate(nn_layers):
        if layer.WhichOneof('layer') == 'custom':
            print( 'layer_id = ', i)
            print( layer)
        else:
            print('{}: layer type: ({}) , inputs: {}, outputs: {}'.
              format(i,layer.WhichOneof('layer'), ", ".join([x for x in layer.input]), ", ".join([x for x in layer.output])))

# We use "netron" for visualization. 
def _visualize(network_path, port_number):
    
    def visualize_using_netron(path, port_number):
        netron.serve_file(path, browse = True, port=port_number)
    
    from threading import Thread
    import time
    
    d = Thread(target = visualize_using_netron, args = (network_path, port_number,))
    d.setDaemon(True)
    d.start()
    time.sleep(5)<jupyter_output><empty_output><jupyter_text>Lets define the first TF graph. This one applies a dense layer and normalizes it. It uses the ["Tile"](https://www.tensorflow.org/versions/master/api_docs/python/tf/tile) op that CoreML does not support.<jupyter_code># define a TF graph: input -> Dense -> unit norm -> output
graph = tf.Graph()
with graph.as_default() as g:
    inputs = tf.placeholder(tf.float32, shape=[None,8], name='input')
    with slim.arg_scope([slim.fully_connected],
          weights_initializer=tf.truncated_normal_initializer(0.0, 0.2),
          weights_regularizer=slim.l2_regularizer(0.0005)):
        y = slim.fully_connected(inputs, 10, scope='fc')
        y = slim.unit_norm(y,dim=1)

output_name = y.op.name
X = np.random.rand(1,8)
frozen_model_file = 'unit_norm_graph.pb'
coreml_model_path = 'unit_norm_graph.mlmodel'
out = _simple_run_and_freeze(graph, output_name, frozen_model_file, feed_dict={'input:0' : X})
print( 'TF out: ', output_name, out.shape, np.sum(out ** 2))
# visualize the frozen TF model
_visualize(frozen_model_file, np.random.randint(8000, 9000))
# Try to convert it : this call should raise an error
coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,8]},
        output_feature_names=['UnitNorm/div:0'])
# we got an unsupported op error. Try again with custom Flag set to true
coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,8]},
        output_feature_names=['UnitNorm/div:0'],
        add_custom_layers=True)<jupyter_output><empty_output><jupyter_text>We can see that the "Tile" op was made into a custom layer in the CoreML model. This op takes in two inputs, it recasts the first one into the shape given by the second input (by repetition). Here is the [documentation](https://www.tensorflow.org/versions/master/api_docs/python/tf/tile).<jupyter_code># visualize CoreML model
_visualize(coreml_model_path, np.random.randint(8000, 9000))<jupyter_output><empty_output><jupyter_text>Note: As we can see in the visualization, the tensors whose values do not change based on the graph inputs (potentially they depend on the shape of the input, which needs to be fixed during conversion) are converted to "load constant" layers in the CoreML graph.<jupyter_code># inspect the CoreML model
spec = coreml_model.get_spec()
_print_coreml_nn_layer_info(spec)<jupyter_output><empty_output><jupyter_text>"ClassName" is an important message: this is the name of the swift/objective-c class that needs to implemented in the Xcode app and will contain the actual code for running the layer.  The "tile" op does not have any parameters, so there is no need to edit generated the coreml specification. Lets now convert a TF graph with the op ["TopKV2"](https://www.tensorflow.org/api_docs/python/tf/nn/top_k) that has parameters.<jupyter_code># define a TF graph: input -> Dense -> softmax -> top_k -> output
tf.reset_default_graph()
graph = tf.Graph()
with graph.as_default() as g:
    x = tf.placeholder(tf.float32, shape=[None,8], name="input")
    y = tf.layers.dense(inputs=x, units=12, activation=tf.nn.relu)
    y = tf.nn.softmax(y, axis=1)
    y = tf.nn.top_k(y, k=3, sorted = False, name='output')
    
output_name = 'output'    
X = np.random.rand(1,8)
frozen_model_file = 'topk_graph.pb'
coreml_model_path = 'topk_graph.mlmodel'
out = _simple_run_and_freeze(graph, output_name, frozen_model_file, feed_dict={'input:0' : X})
print( 'TF out: ', output_name, out.shape, out)
# visualize the frozen TF model
_visualize(frozen_model_file, np.random.randint(8000, 9000))
# Try to convert it : this call should raise an error
coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,8]},
        output_feature_names=['output:0'])
# we got an unsupported op error. Try again with custom Flag set to true
coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,8]},
        output_feature_names=['output:0'],
        add_custom_layers=True)
# inspect the CoreML model
spec = coreml_model.get_spec()
_print_coreml_nn_layer_info(spec)<jupyter_output><empty_output><jupyter_text>[top_k](https://www.tensorflow.org/api_docs/python/tf/nn/top_k) operation has two parameters: 'k' and 'sorted'. In the TF graph, the former is received as an additional input by the op and the latter is an op attribute. Let us modify the MLModel spec directly to add these two parameters to this layer. We need to know a little bit about the custom layer's [proto message](https://github.com/apple/coremltools/blob/5b5b8190764ffe78110be6b4d0edbeebe0253a6e/mlmodel/format/NeuralNetwork.protoL2280) structure to be able to do that.<jupyter_code>nn_layers = coremltools.models.utils._get_nn_layers(spec) # get all the layers as a list
del nn_layers[3].input[1] # delete the second input: its just the value of k
del nn_layers[3].output[1] # delete the second output
nn_layers[3].custom.parameters["k"].intValue = 3
nn_layers[3].custom.parameters["sorted"].boolValue = False
_print_coreml_nn_layer_info(spec)
# save the spec back out
coremltools.models.utils.save_spec(spec, coreml_model_path)
# visualize CoreML model
_visualize(coreml_model_path, np.random.randint(8000, 9000))<jupyter_output><empty_output><jupyter_text>Here is an alternative way to do the same thing using the "custom_conversion_functions" argument:<jupyter_code>def _convert_topk(**kwargs):
    tf_op = kwargs["op"]
    coreml_nn_builder = kwargs["nn_builder"]
    constant_inputs = kwargs["constant_inputs"]
    
    params = NeuralNetwork_pb2.CustomLayerParams()
    params.className = 'Top_K'
    params.description = "Custom layer that corresponds to the top_k TF op"
    params.parameters["sorted"].boolValue = tf_op.get_attr('sorted')
    # get the value of k
    k = constant_inputs.get(tf_op.inputs[1].name, 3)
    params.parameters["k"].intValue = k
    coreml_nn_builder.add_custom(name=tf_op.name,
                                input_names=[tf_op.inputs[0].name],
                                output_names=[tf_op.outputs[0].name],
                                custom_proto_spec=params)

coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,8]},
        output_feature_names=['output:0'],
        add_custom_layers=True,
        custom_conversion_functions={'TopKV2': _convert_topk})

print("\n \n ML Model layers info: \n")
# inspect the CoreML model: this should be same as the one we got above
spec = coreml_model.get_spec()
_print_coreml_nn_layer_info(spec)<jupyter_output><empty_output><jupyter_text>Lets move on to the third and the final example. Now we will encounter an op that is supported but it errors out due to an unsupported coniguration. It is the [Slice](https://www.tensorflow.org/versions/master/api_docs/python/tf/slice) op.<jupyter_code># define a TF graph: input -> conv -> slice -> output
tf.reset_default_graph()
graph = tf.Graph()
with graph.as_default() as g:
    x = tf.placeholder(tf.float32, shape=[None,10,10,3], name="input")
    W = tf.Variable(tf.truncated_normal([1,1,3,5], stddev=0.1))
    y = tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
    y = tf.slice(y, begin=[0,1,1,1], size=[1,2,2,2], name='output')
    
output_name = 'output'    
X = np.random.rand(1,10,10,3)
frozen_model_file = 'slice_graph.pb'
coreml_model_path = 'slice_graph.mlmodel'
out = _simple_run_and_freeze(graph, output_name, frozen_model_file, feed_dict={'input:0' : X})
print( 'TF out: ', output_name, out.shape)
# visualize the frozen TF model
_visualize(frozen_model_file, np.random.randint(8000, 9000))
# Try to convert it : this call should raise an error
coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,10,10,3]},
        output_feature_names=['output:0'])<jupyter_output><empty_output><jupyter_text>This fails, so we provide a custom layer function. Note that this time, the key in the dictionary provided via  "custom_conversion_functions" should be same as the op name ("output")<jupyter_code>def _convert_slice(**kwargs):
    tf_op = kwargs["op"]
    coreml_nn_builder = kwargs["nn_builder"]
    constant_inputs = kwargs["constant_inputs"]
    
    params = NeuralNetwork_pb2.CustomLayerParams()
    params.className = 'Slice'
    params.description = "Custom layer that corresponds to the slice TF op"
    # get the value of begin
    begin = constant_inputs.get(tf_op.inputs[1].name, [0,0,0,0])
    size = constant_inputs.get(tf_op.inputs[2].name, [0,0,0,0])
    # add begin and size as two repeated weight fields
    begin_as_weights = params.weights.add()
    begin_as_weights.floatValue.extend(map(float, begin))
    size_as_weights = params.weights.add()
    size_as_weights.floatValue.extend(map(float, size))
    coreml_nn_builder.add_custom(name=tf_op.name,
                                input_names=[tf_op.inputs[0].name],
                                output_names=[tf_op.outputs[0].name],
                                custom_proto_spec=params)

coreml_model = tfcoreml.convert(
        tf_model_path=frozen_model_file,
        mlmodel_path=coreml_model_path,
        input_name_shape_dict={'input:0':[1,10,10,3]},
        output_feature_names=['output:0'],
        add_custom_layers=True,
        custom_conversion_functions={'output': _convert_slice}) # dictionary has op name as the key

print("\n \n ML Model layers info: \n")
# inspect the CoreML model: this should be same as the one we got above
spec = coreml_model.get_spec()
_print_coreml_nn_layer_info(spec)
# visualize CoreML model
_visualize(coreml_model_path, np.random.randint(8000, 9000))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09924203494347379, 'fraction_numerical': 0.01830678314491264, 'mean_word_length': 4.2615748563703955, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 47, 'https://': 13, 'lorem ipsum': 0, 'www.': 5, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '112426', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 5074, 'n_tokens_neox': 4902, 'n_words': 1343, 'path': 'examples/neural_network_inference/tensorflow_converter/Tensorflow_1/custom_layer_examples.ipynb', 'repo_name': 'DreamChaserMXF/coremltools', 'size': '15568'}