TF-Lite:PReLU 在将 pb 模型转换为 tflite 后工作不正确

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

在pb模型中,我有一个PRelu层。因为 tflite 没有 PRelu OP,所以我将 PRelu 转换为 Relu 像这样:

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

转换为tflite模型时,PRelu将被relu和负OP取代。但是在转换时,两个否定运算之间的否定部分 relu 运算被删除 toco。转换后的模型如下所示:

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

有什么问题吗?

Tensorflow Lite 将激活函数与 op 本身融合在一起,因此 Relu op 将从图中删除。引自 documentation(也提到了 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.


让我们看看这意味着什么。您上面的 TensorFlow 中的 PReLU 代码,使用 TensorBoard 可视化,如下所示(ORIGINAL 图):

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

然而,由于 TfLite 将 Relu 操作与前一个操作融合(更多内容在 docs 中),它将 尝试像这样的东西(注意 [A+B]AB 操作的融合层):

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

但是,由于 neg 操作(一元负)在设计上没有激活函数,ACTUALLY 在 TF-Lite 中发生的事情看起来如下(这是我自己在 1.9.0 版本上测试过):

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

所以,这没有意义!


我个人的解决方法如下(考虑到您已经有一个经过训练的 *.pb 模型并且不想仅仅因为架构发生了变化而重新训练一个新模型):

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)

此代码用于传输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以在离线模式下修改图表。

要简单地解决这个问题,将 PRelu 更改为 最大(x, 0) + alphas* 最小(0, x)