"Dimensions must match" error in tflite conversion with toco

Question

I have a custom CNN model (image classifier) trained with TensorFlow estimator, and I'm gonna use it in iOS app after conversion to TensorFlowLite model.

My model has several dropout layers, and batch normalization layers as well. In order to avoid conversion errors and remove those dropout layers in optimize_for_inference process, I have saved eval_graph.pbtxt separately beside checkpoint files, so as to use it in freeze_graph.

Everything works fine in freeze_graph, and optimize_for_inference throws no errors either. However, after imported both frozen-model and optimized-model files (both .pb) into tensorboard for inspection, I found:

frozen model before optimization

model after optimization

Seems optimize_for_inference removed shapes info of input tensor layer, which is not the case if I freeze the model with graph saved in training mode (the default graph.pbtxt) and optimize it.

Environments:

• Tensorflow 1.8.0 for training;
• Tensorflow 1.13.1 for conversion;

Codes as follow:

Excerpt of model_fn, pretty normal:

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(...)

and custom eval hook 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 and optimize:

# 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 throws error:

# 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

This error seems reasonable since rank 1&2 of shapes are both unknown.

Why?

Answer 1

Yes, it should be after dense:

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