tf2onnx fails with while loop (`KeyError: 'while_body_[num]'`)

[llandsmeer]

A tf.while_loop on model tf.keras.Inputs does not seem to compile. I found out while working a larger code base where an unrolled loop makes both compilation and running very slow, so I tried to make it into a tf.while_loop. Now ONNX conversion fails. The error seems to be largely opset independent.

Minimal test case:

Running

import tf2onnx
import tensorflow as tf

tf.compat.v1.disable_eager_execution()

x0 = tf.keras.Input(16)

def cond(_):
    return True

def body(x):
    return [x]

output, = tf.while_loop(
        cond=cond,
        body=body,
        loop_vars=[x0]
        )

model = tf.keras.Model(inputs=x0, outputs=output)

tf2onnx.convert.from_keras(model, opset=13, output_path='/tmp/test.onnx')

Fails with

Traceback (most recent call last):
  File "/tmp/testcase.py", line 25, in <module>
    tf2onnx.convert.from_keras(model, opset=13, output_path='/tmp/test.onnx')
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tf2onnx/convert.py", line 487, in from_keras
    tf_loader.from_trackable(model, concrete_func, input_names, output_names, large_model)
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tf2onnx/tf_loader.py", line 221, in from_trackable
    frozen_graph = from_function(concrete_func, inputs, outputs, large_model)
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tf2onnx/tf_loader.py", line 273, in from_function
    frozen_func = convert_variables_to_constants_v2(func, lower_control_flow=False, aggressive_inlining=True)
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 1148, in convert_variables_to_constants_v2
    output_graph_def, converted_input_indices = _replace_variables_by_constants(
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 1100, in _replace_variables_by_constants
    input_graph = _GraphDef(converter_data.graph_def)
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 651, in __init__
    self.create_edges()
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 706, in create_edges
    n.create_edges()
  File "/home/llandsmeer/.local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 550, in create_edges
    function = self._enclosing_graph.functions[attr.func.name]
KeyError: 'while_body_3'

self._enclosing_graph.functions seems to be completely empty

Versions

Python 3.10.4
tensorflow==2.9.1
onnx==1.12.0
tf2onnx==1.11.1

I also check with the latest git version, tf2onnx @ git+https://github.com/onnx/tensorflow-onnx@6905d050a3e3fffbcf4a881c13932aaed2bd93be but that still didn't work (same error)

OS

Linux spectre 5.15.0-40-generic #43-Ubuntu SMP Wed Jun 15 12:54:21 UTC 2022 x86_64 x86_64 x86_64 GNU/Linux
Distributor ID: Ubuntu
Description:    Ubuntu 22.04 LTS
Release:        22.04
Codename:       jammy

[llandsmeer]

An error also appears in tensorflow, so maybe it's me not understanding how to use the API in the right way

import tensorflow as tf

tf.compat.v1.disable_eager_execution()

x0 = tf.keras.Input(1)

def cond(_):
    return True

def body(x):
    return [x]

output, = tf.while_loop(
        cond=cond,
        body=body,
        loop_vars=[x0],
        maximum_iterations=1,
        )

model = tf.keras.Model(inputs=x0, outputs=output)

fn = '/tmp/nope'
model.save(fn)
tf.keras.models.load_model(fn)

=>

ValueError: ('There is a cyclic dependency between functions. ....

[llandsmeer]

Anyway, I have function f(x), and I want to calculate f(f(f(f(f(....f(x)))))). There is a Loop construct in ONNX and I have no idea how to access that from tensorflow

[fatcat-z]

Not sure why we need keras here, but there are some examples of how to convert tf.while_loop() to ONNX.

For Keras, there is an example of using tf.while_loop with Keras, please see if you can leverage to construct yours:

_import tensorflow as tf import tf2onnx

class CustomLayer(tf.keras.layers.Layer): def init(self): super(CustomLayer, self).init()

def call(self, inputs):

input_shape = tf.shape(inputs)
end = input_shape[-1]
array = tf.ones((input_shape[-1],))
step = tf.constant(0)
output = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)

def cond(step, output):
    return step < end

def body(step, output):
    output = output.write(step, tf.gather(array, step))
    return step + 1, output

_, final_output = tf.while_loop(cond, body, loop_vars=[step, output])
return tf.reshape(final_output.stack(), (input_shape))

inputs = tf.keras.layers.Input(shape= (None, ), batch_size= 1, name= "timesteps", dtype= tf.int32)

cl = CustomLayer() outputs = cl(inputs) model = tf.keras.Model(inputs, outputs) random_data = tf.random.uniform((1, 7), dtype=tf.int32, maxval=50)

tf2onnx.convert.from_keras(model, opset=13, output_path='/tmp/test.onnx')_

[fatcat-z]

It's been over 3 months, so closing this. Feel free to open a new one if the issue still exists.