Warnings When Converting Graph Containing FusedBatchNormV3 node

[SamG97]

Describe the bug
Trying to convert a model with a FusedBatchNormV3 node causes a lot of warnings. It seems this is because tf2onnx is unable to infer the shape of the node.

Urgency
None

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
• Tensorflow Version: 1.15.0
• Python version: 3.7.3
• tf2onnx version: 1.6.0/55d8ea (latest version on master as of submitting)

To Reproduce
Describe steps/code to reproduce the behavior:

Using the following model:

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import BatchNormalization, Conv2D, Input

x = Input(batch_shape=(1, 1, 32, 32), name="input")
cnn = Conv2D(80, (2, 4), padding="same", name="cnn")(x)
bn = BatchNormalization(name=f"bnorm")(cnn)

model = Model(inputs=x, outputs=bn, name="dummy_model")

I freeze and export the model as a frozen graph. I then run:

python -m tf2onnx.convert \
    --input dummy_model/saved_model.pb \
    --output test_model.onnx \
    --inputs input:0 \
    --outputs bnorm/cond/Merge:0 \
    --opset 11

This does successfully build an ONNX graph which gives the same (minus floating point rounding errors) outputs as the original model. However, it prints a lot of warnings to the console during the process which suggests that something is going wrong:

WARNING:tensorflow:From /home/samg/tensorflow-onnx/tf2onnx/verbose_logging.py:72: The name tf.logging.set_verbosity is deprecated. Please use tf.compat.v1.logging.set_verbosity instead.

2019-11-20 11:01:53,137 - WARNING - From /home/samg/tensorflow-onnx/tf2onnx/verbose_logging.py:72: The name tf.logging.set_verbosity is deprecated. Please use tf.compat.v1.logging.set_verbosity instead.

2019-11-20 11:01:53,215 - INFO - Using tensorflow=1.15.0, onnx=1.6.0, tf2onnx=1.6.0/55d8ea
2019-11-20 11:01:53,215 - INFO - Using opset <onnx, 11>
2019-11-20 11:01:53,215 - WARNING - Cannot infer shape for bnorm/cond/FusedBatchNormV3: bnorm/cond/FusedBatchNormV3:5
2019-11-20 11:01:53,215 - WARNING - Cannot infer shape for bnorm/cond/FusedBatchNormV3_1: bnorm/cond/FusedBatchNormV3_1:5
2019-11-20 11:01:53,238 - INFO - Optimizing ONNX model
2019-11-20 11:01:53,241 - WARNING - Failed to apply optimize_transpose
Traceback (most recent call last):
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/__init__.py", line 52, in optimize_graph
    graph = opt.optimize(current) or graph
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/optimizer_base.py", line 41, in optimize
    graph = self._optimize(graph)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 135, in _optimize
    return self._apply_optimization(graph, self._optimize_at_current_graph_level)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/optimizer_base.py", line 62, in _apply_optimization
    graph = optimize_func(graph)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 148, in _optimize_at_current_graph_level
    if self._handle_nhwc_tranpose(n):
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 269, in _handle_nhwc_tranpose
    return op_handler(trans, p)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 466, in _identity_handler
    self._g.remove_node(node.name)
  File "/home/samg/tensorflow-onnx/tf2onnx/graph.py", line 554, in remove_node
    utils.make_sure(node_name in self._nodes_by_name, "node %s not in current graph, cannot remove", node_name)
  File "/home/samg/tensorflow-onnx/tf2onnx/utils.py", line 291, in make_sure
    raise ValueError("make_sure failure: " + error_msg % args)
ValueError: make_sure failure: node bnorm/cond_Identity__8 not in current graph, cannot remove
2019-11-20 11:01:53,256 - WARNING - Failed to apply optimize_transpose
Traceback (most recent call last):
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/__init__.py", line 52, in optimize_graph
    graph = opt.optimize(current) or graph
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/optimizer_base.py", line 41, in optimize
    graph = self._optimize(graph)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 135, in _optimize
    return self._apply_optimization(graph, self._optimize_at_current_graph_level)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/optimizer_base.py", line 62, in _apply_optimization
    graph = optimize_func(graph)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 148, in _optimize_at_current_graph_level
    if self._handle_nhwc_tranpose(n):
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 269, in _handle_nhwc_tranpose
    return op_handler(trans, p)
  File "/home/samg/tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py", line 415, in _transpose_handler
    self._g.remove_node(node.name)
  File "/home/samg/tensorflow-onnx/tf2onnx/graph.py", line 554, in remove_node
    utils.make_sure(node_name in self._nodes_by_name, "node %s not in current graph, cannot remove", node_name)
  File "/home/samg/tensorflow-onnx/tf2onnx/utils.py", line 291, in make_sure
    raise ValueError("make_sure failure: " + error_msg % args)
ValueError: make_sure failure: node bnorm/cond/FusedBatchNormV3__26 not in current graph, cannot remove
2019-11-20 11:01:53,264 - INFO - After optimization: Const -3 (9->6), Identity -18 (18->0)
2019-11-20 11:01:53,265 - INFO - 
2019-11-20 11:01:53,265 - INFO - Successfully converted TensorFlow model dummy_model/saved_model.pb to ONNX
2019-11-20 11:01:53,266 - INFO - ONNX model is saved at test_model.onnx

Expected behavior
The exceptions causing these warnings shouldn't occur.

Additional context
The frozen model and ONNX graph can be found here.

These warnings only appear when using a model with a tf.keras.layers.BatchNormalization layer in the original definition. In the above example, having the BatchNorm layer as the only layer (i.e. removing the Conv2D layer) will cause the warnings with stack traces during the optimisation phase to go away but still produces the warnings about inferring shapes.

[jignparm]

@SamG97, thanks for the succinct description of the error, and a minimalist way to reproduce it.

The FusedBatchNormV3 model was generating Transpose operators where the output was consumed by ops in a another sub-graph (not the primary graph). The logic to handle that scenario is fixed by PR #747.

Dry run using attached model (link)

...
...
2019-11-22 02:43:22,923 - INFO - Using opset <onnx, 11>
2019-11-22 02:43:22,924 - WARNING - Cannot infer shape for bnorm/cond/FusedBatchNormV3: bnorm/cond/FusedBatchNormV3:5
2019-11-22 02:43:22,924 - WARNING - Cannot infer shape for bnorm/cond/FusedBatchNormV3_1: bnorm/cond/FusedBatchNormV3_1:5
2019-11-22 02:43:22,964 - INFO - Optimizing ONNX model
2019-11-22 02:43:23,007 - INFO - After optimization: Add -1 (1->0), Const -3 (9->6), Identity -17 (17->0), Transpose -3 (6->3)
2019-11-22 02:43:23,008 - INFO - To_ONNX, OK
2019-11-22 02:43:23,010 - INFO - Created d:\tmp\pre-trained\fusedbatchnorm.onnx
2019-11-22 02:43:23,610 - INFO - Model saved to C:\Users\localadmin\AppData\Local\Temp\2\tmp37fn_4qw\run_pretrained\fusedbatchnorm.onnx
2019-11-22 02:43:23,755 - INFO - Run_ONNX OK
2019-11-22 02:43:23,755 - INFO - Results: OK
2019-11-22 02:43:23,758 - INFO - ===================================
2019-11-22 02:43:23,758 - INFO - RESULT: 0 failed of 1, backend=onnxruntime

[SamG97]

Thanks for getting this fixed so fast! I tested the fix on my more complex model that originally caused me to notice the issue, rather than just the simple minimal example above, and that works too now without all the big warnings.