Converting self attention models from TensorFlow to Onnx

[Gerstenberger]

Converting a Transfromer LM from TensorFlow to Onnx and then running the onnx model results in the error

onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from /home/agerstenberger/projects/asr/quantization/onnx/lm/sis-setups/work/apptek_asr/onnx/convert/ConvertTFCheckpointToOnnxJob.41TMCrq7LXoD/output/model.onnx failed:This is an invalid model. 
Type Error: Type 'tensor(bool)' of input parameter (Cast__183:0) of operator (Mul) in node (output/rec/dec_0_self_att_att/MatrixBandPart) is invalid.

using the latest tensorflow and onnx* packages.

One fix could be replacing in the SelfAttentionLayer

returnn/returnn/tf/layers/rec.py

Lines 7997 to 8003 in dbef0ca

	from returnn.tf.util.basic import matrix_triangular
	# (1,1,num_queries,num_keys)
	energy_mask = matrix_triangular((1, 1, num_queries, num_queries), dtype=tf.bool, lower=True)
	if num_queries is not num_keys:
	energy_mask_left = tf.ones((1, 1, num_queries, num_keys - num_queries), dtype=tf.bool)
	energy_mask = tf.concat([energy_mask_left, energy_mask], axis=-1)

using is_onnx_export_global with

...
mask_dtype = tf.bool if not util.is_onnx_export_global() else tf.float32
energy_mask = matrix_triangular((1, 1, num_queries, num_queries), dtype=mask_dtype, lower=True)
if num_queries is not num_keys:
    energy_mask_left = tf.ones((1, 1, num_queries, num_keys - num_queries), dtype=mask_dtype)
    energy_mask = tf.concat([energy_mask_left, energy_mask], axis=-1)
if util.is_onnx_export_global():                                                                                                                                                        
    energy_mask = tf.cast(energy_mask, dtype=tf.bool)
...

This used to work with an older version of tf2onnx. However, the most recent tf2onnx further yields the error during conversion

Tensorflow op [output/rec/dec_0_self_att_att/ones_like: OnesLike] is not supported

and then during runtime

onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from /home/agerstenberger/projects/asr/quantization/onnx/lm/sis-setups/work/apptek_asr/onnx/convert/ConvertTFCheckpointToOnnxJob.41TMCrq7LXoD/output/model.onnx failed:This is an invalid model. 
In Node, ("output/rec/dec_0_self_att_att/ones_like", OnesLike, "", -1) : ("output/rec/dec_0_self_att_att/energy:0": tensor(float),) -> ("output/rec/dec_0_self_att_att/ones_like:0",) , Error No Op registered for OnesLike with domain_version of 15

which is not converted but just copied from the original graph, the default behaviour of tf2onnx for unknown ops (with input tensor output/rec/dec_0_self_att_att/energy:0).
Currently i'm not sure when that changed, but i can review that.

Kind of sounds like a bug to me (have to try out older versions again).
If thats the case, should we wait for a fix? Or provide a workaround?

For a workaround i see two approaches:

1. replace occurences of tf.ones_like in the SelfAttentionLayer with tf.fill ops. TensorFlow actually calls tf.fill in tf.ones, which is called by tf.ones_like. I'm not certain that this is a good idea.
2. provide a custom conversion handler via tf_op decorator (see here) and plug it inside the conversion job. This would also take care of other appearences of tf.ones_like. No need to adapt returnn further.

What are your thoughts?

[albertz]

using the latest tensorflow and onnx* packages.

Can you state which versions exactly?

One fix could be replacing in the SelfAttentionLayer ... using is_onnx_export_global with ... energy_mask = tf.cast(energy_mask, dtype=tf.bool)

I don't understand. Before the change, what dtype is it? I thought it is already bool, or not? So why does this change anything?

This used to work with an older version of tf2onnx

Can you state which version exactly?

And you mean using the cast to bool?

Tensorflow op [output/rec/dec_0_self_att_att/ones_like: OnesLike] is not supported

So, tf.ones_like is not supported in tf2onnx? Is this correct? Don't we use tf.ones_like in many other places as well?

And this worked in a previous tf2onnx version? Did you report this? Or is there an existing issue?

It should be simple to implement our own custom ones_like though, right?

Actually, when I look at the TF code of tf.ones_like of the TF version I have (2.11.0), it also looks like it just calles tf.ones, and that one calls tf.fill. So, maybe in more recent TF versions, they have now a dedicated TF op OnesLike, and use that instead? So then tf2onnx just did not implement this yet? Probably this already has been reported. Please check that and link here. Or otherwise, please report it, and link here. And then let's see what the situation is, when we can expect that this is fixed.

provide a custom conversion handler via tf_op decorator (see here) and plug it inside the conversion job. This would also take care of other appearences of tf.ones_like. No need to adapt returnn further.

tf.ones_like is very common, so I expect that tf2onnx will in any case implement that. The question is when exactly.

Why are you not using an older TF version just for the ONNX export? There should be no problem with that?

[Gerstenberger]

I don't understand. Before the change, what dtype is it? I thought it is already bool, or not? So why does this change anything?

Before the change this would be tf.bool always. However, the implementation of MatrixBandPart (which is tf.compat.v1.matrix_band_part) op in tf2onnx makes use of onnx's Mul operation see which does not support boolean inputs/outputs (for all opsets) see (the same TF op of course does). So i choose one supported data type (in this case float32) for the conversion and then cast the result back to bool.

Maybe it is actually better to adapt the function matrix_triangular instead of the SelfAttentionLayer, to something like

def matrix_triangular(shape, dtype=tf.float32, lower=False, upper=False):                                                                                                                                                                                                
    ...
    assert (lower or upper) and (not lower or not upper)
    x = tf.ones(shape, dtype=dtype)
    onnx_export_bool = util.is_onnx_export_global() and dtype == tf.bool
    if onnx_export_bool:
        # bool type is not supported for tf2onnx op MatrixBandPart 
        x = tf.cast(x, dtype=tf.float32)
    res = tf_compat.v1.matrix_band_part(x, num_lower=-1 if lower else 0, num_upper=-1 if upper else 0)
    if onnx_export_bool:
        return tf.cast(res, dtype=dtype)
    return res

Can you state which version exactly?

I now tried different TF versions. For the environment with

tensorflow-cpu==2.13.1
tf2onnx==1.15.1
onnx==1.15.0
onnxruntime==1.16.1

It works fine.

Using tensorflow-cpu==2.14.0 instead does fail the conversion. There were some changes to tf.ones_like internal implementation it looks like.

There is an onnx operator ConstantOfShape which is equivalent to tf.fill.
Actually, there is a conversion handler already in tf2onnx for ZerosLike, which does basically same thing.

I can create an issue/pr in tf2onnx for that. Better to fix it there.

Why are you not using an older TF version just for the ONNX export? There should be no problem with that?

For now i can use tensorflow==2.13.1 (which is in use currently; will be updated soon i guess).

Thanks!

[Gerstenberger]

Maybe it is actually better to adapt the function matrix_triangular instead of the SelfAttentionLayer, to something like

Or even better would be to fix this in tf2onnx too i guess

[albertz]

Ah I missed that you change the matrix_triangular dtype to mask_dtype. You always needed this, or this is also only with new TF and/or tf2onnx version?

Maybe it is actually better to adapt the function matrix_triangular instead of the SelfAttentionLayer, to something like

Or even better would be to fix this in tf2onnx too i guess

You should also report that (or search for already existing issue), and link here. But I assume this might have not a high priority for them, otherwise they probably would already have fixed that?

I assume the ones_like issue will be more important, as this is very often used. So for that, I would just wait. For matrix_band_part, I think we should keep our current workaround for now.

But your suggestion to move this workaround to matrix_triangular is actually better than having it inside SelfAttentionLayer. So maybe you can do that for now.

And then otherwise keep using TF 2.13 for the export, until it's fixed in tf2onnx. But link the corresponding tf2onnx issue here.

[Gerstenberger]

See the mentioned issues/pr in tf2onnx.

• Unsupported ops TF 2.14.0: OnesLike not supported onnx/tensorflow-onnx#2267
• MatrixBandPart op in tf2onnx does not support bool inputs onnx/tensorflow-onnx#2268

You always needed this, or this is also only with new TF and/or tf2onnx version?

this is independent of the versions.

[Gerstenberger]

The fixes are now merged in tensorflow-onnx.
This is resolved now.