apple / apple/coremltools

Cannot convert keras.layers.MultiHeadAttention

Open
#1,251 8 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

## 🐞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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.