Get ValueError: node *** output needs to be rank 4, is 3 when trying to convert from TF model with Bidirectional ConvLSTM2D Layers

[TylerTheHumanCompiler]

Describe the bug
When trying to convert a tensorflow saved model which contains several Bidirectional ConvLSTM2D Layers with python -m tf2onnx.convert --saved-model ... I get the following error:

ValueError: node BiasAdd output needs to be rank 4, is 3

Urgency
Nope.

System information

• Arch Linux
• 2.2 & 2.5
• Python 3.8

Model
This is the model I used (pastebin.com)

Terminal Output
Output from terminal (pastebin.com)

To Reproduce
Simply use
python -m tf2onnx.convert --saved-model /saved_model_dir --tag=serve --opset 10 --output /output_dir/model.onnx --continue_on_error
from a terminal on a model containing Bidirectional ConvLSTM2D Layers.

Expected behavior
Converting the saved model to onnx without errors.

Additional context
I tried using opset 11 & 12, and using --fold_const as option, but neither worked.

[jignparm]

The code seems to be missing the import headers and also the 'save' function (including tags, signatures etc.)

Can you update with the working code that you are using (to avoid debugging disparity)?

[TylerTheHumanCompiler]

It is updated.

[jignparm]

Thanks for updating. I'm seeing an exception below using the updated script (TF=2.2.0, python=3.7.3, Win-10).

I get the same error, whether or not I train the model (using dummy data).

Does the script run cleanly for you? Wondering if this is a TF version issue.

Train on 1 samples

1/1 [==============================] - ETA: 0s - loss: 1.8556 - output0_loss: 0.4639 - output1_loss: 0.4784 - output2_loss: 0.4391 - output3_loss: 0.4742 - output0_root_mean_squared_error: 0.6811 - output1_root_mean_squared_error: 0.6917 - output2_root_mean_squared_error: 0.6627 - output3_root_mean_squared_error: 0.6886
1/1 [==============================] - 0s 2ms/sample - loss: 1.8556 - output0_loss: 0.4639 - output1_loss: 0.4784 - output2_loss: 0.4391 - output3_loss: 0.4742 - output0_root_mean_squared_error: 0.6811 - output1_root_mean_squared_error: 0.6917 - output2_root_mean_squared_error: 0.6627 - output3_root_mean_squared_error: 0.6886
2020-07-07 23:24:25.078975: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1102] Device interconnect StreamExecutor with strength 1 edge matrix:
2020-07-07 23:24:25.079351: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1108]
Traceback (most recent call last):
  File "train2.py", line 131, in <module>
    model.save(filepath=main_path + '/model/', save_format='tf')
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\engine\network.py", line 1052, in save
    signatures, options)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\save.py", line 138, in save_model
    signatures, options)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\save.py", line 78, in save
    save_lib.save(model, filepath, signatures, options)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\saved_model\save.py", line 951, in save
    obj, export_dir, signatures, options, meta_graph_def)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\saved_model\save.py", line 1008, in _build_meta_graph
    checkpoint_graph_view)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\saved_model\signature_serialization.py", line 75, in find_function_to_export
    functions = saveable_view.list_functions(saveable_view.root)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\saved_model\save.py", line 143, in list_functions
    self._serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\engine\training.py", line 1656, in _list_functions_for_serialization
    Model, self)._list_functions_for_serialization(serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\engine\base_layer_v1.py", line 2439, in _list_functions_for_serialization
    .list_functions_for_serialization(serialization_cache))
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\base_serialization.py", line 87, in list_functions_for_serialization
    fns = self.functions_to_serialize(serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py", line 77, in functions_to_serialize
    serialization_cache).functions_to_serialize)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py", line 92, in _get_serialized_attributes
    serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\model_serialization.py", line 53, in _get_serialized_attributes_internal
    serialization_cache))
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py", line 101, in _get_serialized_attributes_internal
    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\save_impl.py", line 153, in wrap_layer_functions
    original_fns = _replace_child_layer_functions(layer, serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\save_impl.py", line 272, in _replace_child_layer_functions
    serialization_cache).functions)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py", line 92, in _get_serialized_attributes
    serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py", line 101, in _get_serialized_attributes_internal
    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\save_impl.py", line 153, in wrap_layer_functions
    original_fns = _replace_child_layer_functions(layer, serialization_cache)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\save_impl.py", line 272, in _replace_child_layer_functions
    serialization_cache).functions)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py", line 94, in _get_serialized_attributes
    serialized_attr.set_and_validate_objects(object_dict)
  File "D:\anaconda373\lib\site-packages\tensorflow\python\keras\saving\saved_model\serialized_attributes.py", line 212, in set_and_validate_objects
    ' {})'.format(object_dict[key], key))
ValueError: Object dictionary contained a non-trackable object: (None, None) (for key states)

[TylerTheHumanCompiler]

Ah yes, this is a TF2.2 issue. Use either 2.1, 2.3 or higher.
Lost so much time for this bug... 🙄

[jignparm]

@TylerTheHumanCompiler, I can't find TF 2.3 on pypi, but TF2.1 works fine with your script.

The first 3 lines of the code snippet below, placed at the location shown within file nn.py, should fix the error you are seeing.

I'll formalize it as a PR by generalizing it for 1D, 2D and 3D convolutions, and double-checking that the shape alteration doesn't break other scenarios.

tensorflow-onnx/tf2onnx/onnx_opset/nn.py

Line 206 in 9040acc

        if spatial == 2 and len(input_shape) == 3:             # Patch code
            ctx.set_shape(node.input[0], [1] + input_shape)    # Patch code
            input_shape = ctx.get_shape(node.input[0])         # Patch code

        if len(input_shape) != spatial + 2:
            raise ValueError(
                "node {} output needs to be rank {}, is {}".format(
                    node.name, spatial + 2, len(input_shape)
                )
            )

[TylerTheHumanCompiler]

YES!! It totally works! Thanks a lot for the quick response, I really appreciate that!