Custom Keras Layer fails to export
- 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
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