apple / apple/coremltools

Custom Keras Layer fails to export

Open
#1,748 5 comments 0 reactions 0 assignees View on GitHub
bug tf2.x / tf.keras triaged
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug
- Make sure to only create an issue here for bugs in the coremltools Python package. If this is a bug with the Core ML Framework or Xcode, please submit your bug here: https://developer.apple.com/bug-reporting/
- Provide a clear and consise description of the bug.

Custom layer fails to export.

## Stack Trace
- If applicable, please paste the complete stack trace.
```
File "coremltools/coremltools/converters/mil/converter.py", line 190, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
File "coremltools/coremltools/converters/mil/converter.py", line 217, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "coremltools/coremltools/converters/mil/converter.py", line 282, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "coremltools/coremltools/converters/mil/converter.py", line 102, in __call__
return tf2_loader.load()
File "coremltools/coremltools/converters/mil/frontend/tensorflow/load.py", line 61, in load
self._graph_def = self._graph_def_from_model(output_names)
File "coremltools/coremltools/converters/mil/frontend/tensorflow2/load.py", line 133, in _graph_def_from_model
cfs, graph_def = self._get_concrete_functions_and_graph_def()
File "coremltools/coremltools/converters/mil/frontend/tensorflow2/load.py", line 127, in _get_concrete_functions_and_graph_def
graph_def = self._graph_def_from_concrete_fn(cfs)
File "coremltools/coremltools/converters/mil/frontend/tensorflow2/load.py", line 328, in _graph_def_from_concrete_fn
frozen_fn = _convert_variables_to_constants_v2(cfs[0], lower_control_flow=False, aggressive_inlining=True)
File ".local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 1162, in convert_variables_to_constants_v2
return _construct_concrete_function(func, output_graph_def,
File ".local/lib/python3.10/site-packages/tensorflow/python/framework/convert_to_constants.py", line 1087, in _construct_concrete_function
new_func = wrap_function.function_from_graph_def(output_graph_def,
File ".local/lib/python3.10/site-packages/tensorflow/python/eager/wrap_function.py", line 655, in function_from_graph_def
wrapped_import = wrap_function(_imports_graph_def, [])
File ".local/lib/python3.10/site-packages/tensorflow/python/eager/wrap_function.py", line 619, in wrap_function
func_graph.func_graph_from_py_func(
File ".local/lib/python3.10/site-packages/tensorflow/python/framework/func_graph.py", line 1247, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File ".local/lib/python3.10/site-packages/tensorflow/python/eager/wrap_function.py", line 83, in __call__
return self.call_with_variable_creator_scope(self._fn)(*args, **kwargs)
File ".local/lib/python3.10/site-packages/tensorflow/python/eager/wrap_function.py", line 89, in wrapped
return fn(*args, **kwargs)
File ".local/lib/python3.10/site-packages/tensorflow/python/eager/wrap_function.py", line 649, in _imports_graph_def
importer.import_graph_def(graph_def, name="")
File ".local/lib/python3.10/site-packages/tensorflow/python/util/deprecation.py", line 561, in new_func
return func(*args, **kwargs)
File ".local/lib/python3.10/site-packages/tensorflow/python/framework/importer.py", line 403, in import_graph_def
return _import_graph_def_internal(
File ".local/lib/python3.10/site-packages/tensorflow/python/framework/importer.py", line 505, in _import_graph_def_internal
raise ValueError(str(e))
ValueError: Node 'model/multi_head_positional_embedding/Gather/axis' is not unique
```

## To Reproduce
- Please add a minimal code example that can reproduce the error when running it.
```python
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, models

class MultiHeadPositionalEmbedding(keras.layers.Layer):
def __init__(self, query_height=-1, key_height=-1, **kwargs):
super(MultiHeadPositionalEmbedding, self).__init__(**kwargs)
self.query_height, self.key_height = query_height, key_height

def build(self, input_shape, **kwargs):
_, num_heads, qq_blocks, kk_blocks = input_shape
self.bb = self.add_weight(name="positional_embedding", shape=(kk_blocks, num_heads), initializer="zeros", trainable=True)

if self.query_height == -1:
q_blocks_h = q_blocks_w = int(tf.math.sqrt(float(qq_blocks))) # hh == ww
else:
q_blocks_h, q_blocks_w = self.query_height, int(qq_blocks / self.query_height)

strides = int(tf.math.ceil(tf.math.sqrt(float(kk_blocks / qq_blocks))))
if self.key_height == -1:
k_blocks_h = q_blocks_h * strides
while kk_blocks % k_blocks_h != 0:
k_blocks_h -= 1
k_blocks_w = int(kk_blocks / k_blocks_h)
else:
k_blocks_h, k_blocks_w = self.key_height, int(kk_blocks / self.key_height)
self.k_blocks_h, self.k_blocks_w = k_blocks_h, k_blocks_w
# print(f"{q_blocks_h = }, {q_blocks_w = }, {k_blocks_h = }, {k_blocks_w = }, {strides = }")

x1, y1 = tf.meshgrid(range(q_blocks_h), range(q_blocks_w))
x2, y2 = tf.meshgrid(range(k_blocks_h), range(k_blocks_w))
aa = tf.concat([tf.reshape(x1, (-1, 1)), tf.reshape(y1, (-1, 1))], axis=-1)
bb = tf.concat([tf.reshape(x2, (-1, 1)), tf.reshape(y2, (-1, 1))], axis=-1)
# print(f">>>> {aa.shape = }, {bb.shape = }") # aa.shape = (16, 2), bb.shape = (49, 2)
cc = [tf.math.abs(bb - ii * strides) for ii in aa]
self.bb_pos = tf.stack([ii[:, 0] + ii[:, 1] * k_blocks_h for ii in cc], name="my_stack")

# print(f">>>> {self.bb_pos.shape = }") # self.bb_pos.shape = (16, 49)

super(MultiHeadPositionalEmbedding, self).build(input_shape)

def call(self, inputs, **kwargs):
pos_bias = tf.gather(self.bb, self.bb_pos)
pos_bias = tf.transpose(pos_bias, [2, 0, 1])
return inputs + pos_bias

def get_config(self):
base_config = super().get_config()
base_config.update({"query_height": self.query_height, "key_height": self.key_height})
return base_config

if __name__ == "__main__":
import coremltools as ct
import numpy as np
t = layers.Input(
shape=(8, 64, 256)
)
x = MultiHeadPositionalEmbedding()(t)

m = models.Model([t], [x])
q = m.predict(np.zeros((1,8, 64, 256)))
print(m.summary(), q)

coreml_model = ct.convert(m,
minimum_deployment_target=ct.target.iOS16,
inputs=[],
source='tensorflow')
```

```
Model: "model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_1 (InputLayer) [(None, 8, 64, 256)] 0

multi_head_positional_embed (None, 8, 64, 256) 2048
ding (MultiHeadPositionalEm
bedding)

=================================================================
Total params: 2,048
Trainable params: 2,048
Non-trainable params: 0
_________________________________________________________________
None [[[[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]

[[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]

[[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]

...

[[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]

[[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]

[[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]]]

```

## System environment (please complete the following information):
- coremltools version:
- OS Linux Ubuntu 22.04
- Tensorflow 2.10:
- Coremltools 6.1

Contributor guide

Open the contributing guide

Research direction

Reproduce the failure with the provided TensorFlow 2.10 example and inspect the conversion path in coremltools/converters/mil/frontend/tensorflow2/load.py, especially _graph_def_from_concrete_fn, along with converter.py. Trace why freezing the custom layer produces a non-unique Gather/axis node. Done means the supplied custom Keras layer exports successfully without the ValueError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.