apple / apple/coremltools

Tensorflow/Keras mixed_float16 export model fails

Open
#1,137 3 comments 0 reactions 0 assignees View on GitHub
bug tf2.x / tf.keras
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.