Tensorflow/Keras mixed_float16 export model fails
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🐞Describe the bug
- Using TensorFlow/keras with mixed precision training fails to export
- Keras converter issue
NotImplementedError: Cast: Provided destination type fp16 not supported.
## To Reproduce
- If a python script can reproduce the error, please paste the code snippet
```
import os
import tensorflow as tf
import numpy as np
from tensorflow.keras import mixed_precision
import coremltools as ct
def mnist_dataset(batch_size):
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
# The `x` arrays are in uint8 and have values in the range [0, 255].
# You need to convert them to float32 with values in the range [0, 1]
x_train = x_train / np.float32(255)
y_train = y_train.astype(np.int64)
train_dataset = tf.data.Dataset.from_tensor_slices(
(x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
return train_dataset
def build_and_compile_cnn_model():
model = tf.keras.Sequential([
tf.keras.Input(shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10, dtype='float32')
])
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
metrics=['accuracy'])
print(model.output)
return model
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)
print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)
batch_size = 64
single_worker_dataset = mnist_dataset(batch_size)
single_worker_model = build_and_compile_cnn_model()
single_worker_model.fit(single_worker_dataset, epochs=3, steps_per_epoch=70)
single_worker_model.save('tf_keras_model')
mlmodel = ct.convert('tf_keras_model')
mlmodel.save("test.mlmodel")
```
## System environment (please complete the following information):
- coremltools version : 4.1
- OS : Ubuntu 20.04
- How you install python (anaconda, virtualenv, system): system
- python version (e.g. 3.7): 3.8.5
- any other relevant information: TensorFlow 2.4.1
Contributor guide
Research direction
Start by running the provided TensorFlow/Keras mixed_float16 reproduction with coremltools 4.1 and TensorFlow 2.4.1, focusing on the ct.convert('tf_keras_model') entry point and the reported Cast fp16 error. Trace the converter path handling mixed-precision model types. Done means the example exports and saves test.mlmodel without the NotImplementedError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- keras, python, tensorflow
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100