"Dimensions must match" 使用 toco 进行 tflite 转换时出错

Question

我有一个使用 TensorFlow 估算器训练的自定义 CNN 模型（图片 classifier），我将在转换为 TensorFlowLite 模型后在 iOS 应用程序中使用它。

我的模型有几个 dropout 层，还有 batch normalization 层。为了避免转换错误并删除 optimize_for_inference 过程中的那些丢失层，我在检查点文件旁边单独保存了eval_graph.pbtxt，以便在freeze_graph中使用它。

在 freeze_graph 中一切正常，optimize_for_inference 也没有抛出任何错误。然而，将冻结模型和优化模型文件（均为.pb）导入tensorboard进行检查后，我发现：

优化前冻结模型

优化后的模型

似乎optimize_for_inference 删除了输入张量层的形状信息，如果我用训练模式（默认graph.pbtxt）保存的图形冻结模型并优化它，情况就不是这样了。

环境：

用于训练的 Tensorflow 1.8.0；
用于转换的 Tensorflow 1.13.1；

代码如下：

model_fn 的摘录，非常正常：

def cnn_model_fn(features, labels, mode, params):
    """Model function for CNN."""
    # Input Layer, images aleady reshaped before feed in;
    net = tf.placeholder_with_default(
        features['Pixels'],
        (None, 48, 48, 1),
        name='input_tensor'
    )

    # bn-1
    net = tf.layers.batch_normalization(
        inputs=net,
        training=mode == tf.estimator.ModeKeys.TRAIN
    )

    # conv2d-1
    net = tf.layers.conv2d(
        inputs=net,
        filters=32,
        kernel_size=[3, 3],
        padding='same',
        activation=tf.nn.relu
    )

    # conv2ds, dropouts, poolings, bns...

    # CONV2Ds -> DENSEs
    # 48 pixels pooled three times (kernel_sizes=2, strides=2), and final conv2d has 128 neurons;
    net = tf.reshape(net, [-1, 6 * 6 * 128])

    # bn-4
    net = tf.layers.batch_normalization(
        inputs=net,
        training=mode == tf.estimator.ModeKeys.TRAIN
    )

    # dense-1
    net = tf.layers.dense(
        inputs=net,
        units=256,
        kernel_regularizer=keras.regularizers.l2(0.001),
        activation=tf.nn.relu
    )

    # denses, logits, nothing special...

    # In prediction:
    if mode == tf.estimator.ModeKeys.PREDICT:        
        return tf.estimator.EstimatorSpec(...)

    # In evaluation:
    if mode == tf.estimator.ModeKeys.EVAL:
        # hook for saving graph in eval mode, this graph will be used in freezing & optimizing process;
        eval_finish_hook = EvalFinishHook()
        eval_finish_hook.model_dir = params['model_dir']
        return tf.estimator.EstimatorSpec(
            ...,
            evaluation_hooks=[eval_finish_hook]
        )

    # In training:
    if mode == tf.estimator.ModeKeys.TRAIN:
        return tf.estimator.EstimatorSpec(...)

和自定义评估挂钩 class:

class EvalFinishHook(tf.train.SessionRunHook):
    model_dir = '.'
    _saver = None

    def begin(self):
        self._saver = tf.train.Saver()
        super().begin()

    def end(self, session):
        dst_dir = self.model_dir + 'eval_ckpt'
        self._saver.save(sess=session, save_path=dst_dir + '/eval.ckpt')
        tf.train.write_graph(session.graph.as_graph_def(), dst_dir, 'eval_graph.pbtxt')
        super().end(session)

冻结并优化：

# freeze graph
echo "freezing checkpoint ${best_step}..."
freeze_graph \
--input_graph=${input_graph} \
--input_checkpoint=${input_checkpoint} \
--input_binary=false \
--output_graph=${frozen_model} \
--output_node_names=${output_names} \

# optimize for inference
echo "optimizing..."
/path/to/bazel-bin/tensorflow/python/tools/optimize_for_inference \
--input=${frozen_model} \
--output=${optimized_model} \
--frozen_graph=True \
--input_names=${input_names} \
--output_names=${output_names}

toco 抛出错误：

# convert to tflite
echo "converting..."
toco \
--graph_def_file=${optimized_model} \
--input_format=TENSORFLOW_GRAPHDEF \
--output_format=TFLITE \
--inference_type=FLOAT \
--input_type=FLOAT \
--input_arrays=${input_names} \
--output_arrays=${output_names} \
--input_shapes=1,48,48,1 \
--output_file=${tflite_model}


# error info
Check failed: dim_x == dim_y (128 vs. 4608)Dimensions must match

这个错误似乎是合理的，因为形状的等级 1 和 2 都是未知的。

为什么？

Answer 1

optimize_for_inference 从图中随机删除 dropout 层，通常在输入上使用 dropout。因此答案可能是肯定的。

bazel-bin/tensorflow/python/tools/optimize_for_inference \ 
--input=/tf_files/retrained_graph.pb \ 
--output=/tf_files/optimized_graph.pb \ 
--input_names={} \ 
--output_names=result

让我们尝试使用 RandomUniform、FLOOR、TensorFlowShape、TensorFlowSwitch、TensorFlowMerge 自定义实现，以禁用错误。

参考：Dropout Regularization

Answer 2

当您使用冻结图而不是 graph.pbtxt 时，您应该使用 eval.pbtxt。

请检查tensorflow/models

所以让我们将第一个维度的 'None' 替换为零，其余维度替换为描述向量/矩阵的大小。另一点是观察矩阵乘法规则，即第一个操作数的列数必须与第二个操作数的行数匹配。

如果对你有帮助，请采纳

Answer 3

好吧，似乎交换 bn-4 和 dense-1 可以消除错误。因此，在这种情况下，批量归一化应该在密集之后进行（比如，在 conv2d->dense reshape 之后）。

Answer 4

对，应该在dense之后：

model.add(Dense(.., ..))
model.add(BatchNormalization())
model.add(Activation(...))
model.add(Dropout(...))

"Dimensions must match" 使用 toco 进行 tflite 转换时出错

"Dimensions must match" error in tflite conversion with toco

python

tensorflow

tensorboard

tensorflow-estimator

tensorflow-lite