TF while loop error

[marton-avrios]

I exported the original T5 model (https://github.com/google-research/text-to-text-transfer-transformer) into saved model format and then tried to convert it with onnx (python -m tf2onnx.convert --saved-model my_saved_model_dir --output model.onnx --opset 13 --tag serve --signature_def serving_default). I got the following error (TF 2.3.0, python 3.7):

ERROR - rewriter <function rewrite_generic_loop at 0x7f4458e8cae8>: exception 'while_loop/while/Enter_1'
Traceback (most recent call last):
  File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/convert.py", line 488, in <module>
    main()
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/convert.py", line 256, in main
    output_path=args.output)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/convert.py", line 153, in _convert_common
    g = process_tf_graph(tf_graph, const_node_values=const_node_values, **kwargs)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 534, in process_tf_graph
    output_names, initialized_tables, outputs_to_values, outputs_to_dtypes, op_cnt, attr_cnt)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 628, in process_parsed_graph
    run_rewriters(g, rewriters, continue_on_error)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 363, in run_rewriters
    raise ex
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 354, in run_rewriters
    ops = func(g, g.get_nodes())
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/rnn.py", line 50, in rewrite_generic_loop
    return LoopRewriter(g).run()
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/loop_rewriter.py", line 36, in run
    return self.run_internal()
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/loop_rewriter_base.py", line 213, in run_internal
    context.cond_graph = self._crop_loop_condition_sub_graph(context)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/loop_rewriter_base.py", line 295, in _crop_loop_condition_sub_graph
    loop_var = context.loop_properties.all_variables[enter_node.name]
KeyError: 'while_loop/while/Enter_1'

Here is the full output:

/usr/lib/python3.7/runpy.py:125: RuntimeWarning: 'tf2onnx.convert' found in sys.modules after import of package 'tf2onnx', but prior to execution of 'tf2onnx.convert'; this may result in unpredictable behaviour
  warn(RuntimeWarning(msg))
2021-06-10 11:58:42,927 - INFO - Signatures found in model: [serving_default].
2021-06-10 11:58:42,927 - INFO - Output names: ['inputs', 'inputs_encoded', 'outputs', 'scores']
WARNING:tensorflow:Issue encountered when serializing global_step.
Type is unsupported, or the types of the items don't match field type in CollectionDef. Note this is a warning and probably safe to ignore.
to_proto not supported in EAGER mode.
2021-06-10 11:58:43,096 - WARNING - Issue encountered when serializing global_step.
Type is unsupported, or the types of the items don't match field type in CollectionDef. Note this is a warning and probably safe to ignore.
to_proto not supported in EAGER mode.
WARNING:tensorflow:Issue encountered when serializing variables.
Type is unsupported, or the types of the items don't match field type in CollectionDef. Note this is a warning and probably safe to ignore.
to_proto not supported in EAGER mode.
2021-06-10 11:58:43,106 - WARNING - Issue encountered when serializing variables.
Type is unsupported, or the types of the items don't match field type in CollectionDef. Note this is a warning and probably safe to ignore.
to_proto not supported in EAGER mode.
WARNING:tensorflow:Issue encountered when serializing trainable_variables.
Type is unsupported, or the types of the items don't match field type in CollectionDef. Note this is a warning and probably safe to ignore.
to_proto not supported in EAGER mode.
2021-06-10 11:58:43,107 - WARNING - Issue encountered when serializing trainable_variables.
Type is unsupported, or the types of the items don't match field type in CollectionDef. Note this is a warning and probably safe to ignore.
to_proto not supported in EAGER mode.
WARNING:tensorflow:From /home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tf_loader.py:603: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.compat.v1.graph_util.extract_sub_graph`
2021-06-10 11:59:12,199 - WARNING - From /home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tf_loader.py:603: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.compat.v1.graph_util.extract_sub_graph`
2021-06-10 11:59:16,967 - INFO - Using tensorflow=2.3.0, onnx=1.9.0, tf2onnx=1.8.5/50049d
2021-06-10 11:59:16,967 - INFO - Using opset <onnx, 13>
2021-06-10 11:59:21,494 - INFO - Computed 12 values for constant folding
2021-06-10 11:59:27,376 - WARNING - shape missmatch between caller and function: __inference_Dataset_map_lambda_157
2021-06-10 11:59:27,435 - INFO - Computed 0 values for constant folding
2021-06-10 11:59:27,767 - INFO - folding node using tf type=OneHot, name=decoder/block_000/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:27,788 - INFO - folding node using tf type=OneHot, name=decoder/block_001/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:27,809 - INFO - folding node using tf type=OneHot, name=decoder/block_002/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:27,831 - INFO - folding node using tf type=OneHot, name=decoder/block_003/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:27,853 - INFO - folding node using tf type=OneHot, name=decoder/block_004/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:27,877 - INFO - folding node using tf type=NextIteration, name=while_loop/while/NextIteration_1
2021-06-10 11:59:27,878 - INFO - folding node using tf type=OneHot, name=while_loop/while/encoder/block_000/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:28,001 - INFO - folding node using tf type=OneHot, name=while_loop/while/encoder/block_001/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:28,122 - INFO - folding node using tf type=OneHot, name=while_loop/while/encoder/block_002/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:28,237 - INFO - folding node using tf type=OneHot, name=while_loop/while/encoder/block_003/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:28,359 - INFO - folding node using tf type=OneHot, name=while_loop/while/encoder/block_004/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:28,477 - INFO - folding node using tf type=OneHot, name=while_loop/while/encoder/block_005/layer_000/SelfAttention/one_hot/parallel_0/one_hot
2021-06-10 11:59:31,141 - ERROR - rewriter <function rewrite_generic_loop at 0x7ffa1103fa60>: exception 'while_loop/while/Enter_1'
Traceback (most recent call last):
  File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/convert.py", line 488, in <module>
    main()
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/convert.py", line 256, in main
    output_path=args.output)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/convert.py", line 153, in _convert_common
    g = process_tf_graph(tf_graph, const_node_values=const_node_values, **kwargs)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 534, in process_tf_graph
    output_names, initialized_tables, outputs_to_values, outputs_to_dtypes, op_cnt, attr_cnt)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 628, in process_parsed_graph
    run_rewriters(g, rewriters, continue_on_error)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 363, in run_rewriters
    raise ex
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/tfonnx.py", line 354, in run_rewriters
    ops = func(g, g.get_nodes())
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/rnn.py", line 50, in rewrite_generic_loop
    return LoopRewriter(g).run()
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/loop_rewriter.py", line 36, in run
    return self.run_internal()
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/loop_rewriter_base.py", line 213, in run_internal
    context.cond_graph = self._crop_loop_condition_sub_graph(context)
  File "/home/martongyorgy/myenv/lib/python3.7/site-packages/tf2onnx/rewriter/loop_rewriter_base.py", line 295, in _crop_loop_condition_sub_graph
    loop_var = context.loop_properties.all_variables[enter_node.name]
KeyError: 'while_loop/while/Enter_1'

[TomWildenhain-Microsoft]

I wasn't able to get the pretrained checkpoints loaded to reproduce this, but taking a shot in the dark it seems folding the NextIteration op might be causing the issue. Try it again with #1561. If that fails, it would be very helpful to have a zipped copy of your saved model. It doesn't need to be trained, just to test conversion.

[marton-avrios]

The problem is with a custom while loop I added to the original model. However it is still strange, that tf has no problem with it while ONNX has. Unfortunately it is a proprietary model so I cannot add example here but will try to open another issue with reproducability later.

[guschmue]

in general we should be able deal with tf.while but the old tensorflow control flow requires us to pattern match in the graph and maybe there is something in your model where this goes wrong.