[FIXED] TF-Lite: PReLU works incorrectly after converting pb model to tflite

Issue

In the pb model, I have a PRelu layer. Because tflite do not have PRelu OP, so I convert the PRelu to Relu like this:

 pos = relu(x)
 neg = - alphas * relu(-x)
 return pos + neg

When converting to tflite model, the PRelu will be replaced by relu and negative OPs. But when converting, the neg part relu op between two negative op is dropped by toco. The converted model looked like this:

pos = relu(x)
neg = -alphas * (-x)
return pos + neg

any problem?

Solution

Tensorflow Lite fuses the activation function with the op itself, so Relu ops will be removed from the graph. Quote from documentation (that also mentions about tf.nn.relu):

Note that many of those operations don't have TensorFlow Lite equivalents and the corresponding model will not be convertible if they can't be elided or fused.

---

Let's see what it means. Your above code of PReLU in TensorFlow, visualized using TensorBoard, looks like following (the ORIGINAL graph):

conv --> relu ---------------------\
     \-> neg -> relu -> mul -> neg --> add

However, due to the fact that TfLite fuses the Relu ops with the previous op (more in docs), it will TRY to make something like this (note that [A+B] is a fused layer of A and B ops):

[conv+relu] -----------------------------\
            \-> [neg+relu] -> mul -> neg --> add

But, since neg operation (unary minus) doesn't have activation function by design, what ACTUALLY happens inside TF-Lite looks like following (this was tested by myself on version 1.9.0):

[conv+relu] ----------------------\
            \-> neg -> mul -> neg --> add

So, it makes no sense!

---

My personal work-around of this is the following (considering you already have a trained *.pb model and don't want to re-train a new model just because the architecture has changed):

def tflite_tolerant_prelu(_x, alpha, name_scope):
    with tf.name_scope(name_scope):
        alpha = tf.constant(alpha, name='alpha')
        return tf.maximum(_x, 0) + alpha * tf.minimum(_x, 0)

def replace_prelu(graph, prelu_block_name, tensor_before, tensor_after):
    alpha = graph.get_tensor_by_name(os.path.join(prelu_block_name, 'alpha:0'))
    with tf.Session() as sess:
        alpha_val = alpha.eval()
    new_prelu = tflite_tolerant_prelu(tensor_before,
            alpha_val, prelu_block_name + '_new')
    tf.contrib.graph_editor.swap_inputs(tensor_after.op, [new_prelu])

before = mtcnn_graph.get_tensor_by_name('pnet/conv1/BiasAdd:0')
after = mtcnn_graph.get_tensor_by_name('pnet/pool1:0')
replace_prelu(mtcnn_graph, 'pnet/PReLU1', before, after)

This code was used to transfer MTCNN from TensorFlow to TensorFlow Lite. Looks a bit ugly (definitely needs make it look cleaner), but it is fully functional and does the job. Note that I used the graph editor tool tensorflow.contrib.graph_editor to modify the graph in offline-mode.

Answered By - hav4ik