CODEWITHSUNDEEP
pythontensorflowtensorboardtensorflow-estimatortensorflow-lite
Ke Yang
Ke Yang

Reputation: 921

"Dimensions must match" error in tflite conversion with toco

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 freezed model before optimization

model after 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:

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?

Upvotes: 2

Views: 1988

Answers (4)

Ke Yang
Ke Yang

Reputation: 921

well, it seems that swap bn-4 and dense-1 mutes the error. So batch normalization should come AFTER dense in this occasion(say, right behind conv2d->dense reshape).

Upvotes: 1

Mahsa Hassankashi
Mahsa Hassankashi

Reputation: 2139

Yes, it should be after dense:

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

Upvotes: 3

Mahsa Hassankashi
Mahsa Hassankashi

Reputation: 2139

You should use eval.pbtxt when you are using froze graph instead of graph.pbtxt.

Please let check tensorflow/models

So let`s replace 'None' as zero for the first dimension and for the rest ones replaced size of description vector / matrix. Another point is to observe matrix multiplication rule which says number of columns of the first operand must match the number of rows of the second operand.

If it helps you, please accept answer.

Upvotes: 0

Mahsa Hassankashi
Mahsa Hassankashi

Reputation: 2139

optimize_for_inference removes dropout layers from the graph as randomly, it is common to use dropout on the inputs. Therefore the answer could be YES.

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

Let's try to custom implementation with RandomUniform, FLOOR, TensorFlowShape, TensorFlowSwitch, TensorFlowMerge, in order to disable error.

Reference: Dropout Regularization

Upvotes: 2

Related Questions