Cannot convert keras.layers.MultiHeadAttention
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🐞Describe the bug
There is an issue when converting TF2 Keras models that contain MultiHeadAttention:
```py
layers.MultiHeadAttention(num_heads=num_heads, key_dim=head_dim, name="attention")
```
The conversion fails with: `ValueError: Cannot add const [512*is10, 512]`
The `is10` variable increments each time I try.
The problem seems to be when calculating the matrix size for one of the Einsums. I can't tell if it's the Q,K,V einsums causing trouble or the scaled dot product einsums.
I tried also using the MultiHeadAttention from TensorFlow addons, but that one failed with unsupported einsums.
The model trains and executes fine so this seems to be a conversion issue. I tried 5.0b2.
## Trace
```py
coremltools/converters/mil/frontend/tensorflow/ops.py in Einsum(context, node)
493 a = context[node.inputs[0]]
494 b = context[node.inputs[1]]
--> 495 x = build_einsum_mil(a, b, equation, node.name)
496 context.add(node.name, x)
497
coremltools/converters/mil/frontend/_utils.py in build_einsum_mil(a_var, b_var, equation, name)
66 if parsed_vectors_rev == ([0,1,2],[2,3,4],[0,1,3,4]):
67 a_var, b_var = _swap(a_var, b_var)
---> 68 x_1 = mb.reshape(x=a_var, shape=[a_var.shape[0] * a_var.shape[1], a_var.shape[2]])
69 x_2 = mb.reshape(x=b_var, shape=[b_var.shape[0], b_var.shape[1] * b_var.shape[2]])
70 x = mb.matmul(x=x_1, y=x_2, transpose_x=False, transpose_y=False)
coremltools/converters/mil/mil/ops/registry.py in add_op(cls, **kwargs)
59 @classmethod
60 def add_op(cls, **kwargs):
---> 61 return cls._add_op(op_cls, **kwargs)
62
63 setattr(Builder, op_type, add_op)
coremltools/converters/mil/mil/builder.py in _add_op(cls, op_cls, **kwargs)
160 # Shallow copy list inputs to ensure op inputs are immutable
161 kwargs = {k: v if not isinstance(v, (list, tuple)) else v[:] for k, v in kwargs.items() if v is not None}
--> 162 kwargs.update(cls._create_vars(
163 input_spec=op_cls.input_spec,
164 op_name=kwargs["name"], before_op=before_op,
coremltools/converters/mil/mil/builder.py in _create_vars(cls, input_spec, op_name, before_op, candidate_kv)
143 if isinstance(in_type, (ScalarOrTensorInputType,
144 ListOrScalarOrTensorInputType)):
--> 145 var = cls._add_const(val, new_var_name, before_op)
146 update_dict[k] = var
147
~/.virtualenvs/codepredictor/lib/python3.9/site-packages/coremltools/converters/mil/mil/builder.py in _add_const(cls, val, name, before_op)
73 def _add_const(cls, val, name, before_op):
74 if not is_python_value(val):
---> 75 raise ValueError("Cannot add const {}".format(val))
76 if any_symbolic(val):
77 msg = (
ValueError: Cannot add const [512*is10, 512]
```
## To Reproduce
The model source here reproduces this bug: https://github.com/keras-team/keras-io/blob/master/examples/generative/text_generation_with_miniature_gpt.py
## System environment:
- coremltools version: 5.0b2
- OS: MacOS
- macOS version: 11.4
- XCode version: 12.5.1
- How you install python: virtualenv
- python version: 3.9.4
- TensorFlow: 2.5
Contributor guide
Research direction
Start with the Einsum handler in coremltools/converters/mil/frontend/tensorflow/ops.py and build_einsum_mil in coremltools/converters/mil/frontend/_utils.py, where the symbolic reshape fails. Run the Keras miniature GPT example linked in the issue with the reported TensorFlow and coremltools versions. Done means the MultiHeadAttention model converts without the symbolic-constant ValueError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, tensorflow
- Domain
- devtools, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100