Unified Conversion API

With coremltools 4.0 and newer versions, use the Unified Conversion API to convert to Core ML models from the following source frameworks:

• TensorFlow 1
• TensorFlow 2
• TensorFlow's Keras APIs
• PyTorch

🚧

API compatibility

The Unified Conversion API supports only TensorFlow and PyTorch neural networks. For converting models using multi-backend Keras, Caffe, and ONNX, use the conversion APIs specific to those packages.

Supported formats

Formats supported by the Unified Conversion API include the following:

Framework	Supported Format
TensorFlow versions 1.x	- Frozen tf.Graph - Frozen graph (.pb) file path - tf.keras.Model - HDF5 file path (.h5) - SavedModel directory path
TensorFlow versions 2.x	- tf.keras.Model - HDF5 file path (.h5) - SavedModel directory path - A [concrete function] (https://www.tensorflow.org/guide/concrete_function)
PyTorch	- TorchScript object - TorchScript object saved as a .pt file

Load and convert a model

The typical conversion process with the Unified Conversion API is to load the model to infer its type, and then use the convert() method to convert it to the Core ML format. Follow these steps:

1. Load the model.

TensorFlow

import coremltools as ct

# Load TensorFlow model
import tensorflow as tf # Tf 2.2.0

tf_model = tf.keras.applications.MobileNet()

PyTorch

import coremltools as ct

# Load PyTorch model (and perform tracing)
torch_model = torchvision.models.mobilenet_v2()
torch_model.eval() 

example_input = torch.rand(1, 3, 256, 256)
traced_model = torch.jit.trace(torch_model, example_input)

2. Convert the model using convert():

TensorFlow

# Convert using the same API
model_from_tf = ct.convert(tf_model)

PyTorch

# Convert using the same API. Note that we need to provide "inputs" for pytorch conversion.
model_from_torch = ct.convert(traced_model,
                              inputs=[ct.TensorType(name="input", shape=example_input.shape)])

The conversion produces an MLModel object which you can use to make predictions, change metadata, or save to the Core ML format for use in Xcode. For more information, see the Core ML Overview.

📘

Conversion options

The convert() method tries to infer as much as possible from the source network, but some information may not be present, such as input names, types, shapes, and classifier options. For more information see Conversion Options.

Convert from TensorFlow 2

TensorFlow 2 models are typically exported as tf.Model objects in the SavedModel or HDF5 file formats. For additional TensorFlow formats you can convert, see the TensorFlow 2 Conversion page.

The following example demonstrates how to use the convert() method to convert an Xception model from tf.keras.applications:

Python

import coremltools as ct 
import tensorflow as tf

# Load from .h5 file
tf_model = tf.keras.applications.Xception(weights="imagenet", 
                                          input_shape=(299, 299, 3))

# Convert to Core ML
model = ct.convert(tf_model)

Convert from TensorFlow 1

The conversion API can also convert models from TensorFlow 1. These models are generally exported with the extension .pb, in the frozen protobuf file format, using TensorFlow 1's freeze graph utility. You can pass this model directly into the convert() method. For details, see the TensorFlow 1 Conversion page.

The following example demonstrates how to convert a pre-trained MobileNet model in the frozen protobuf format to Core ML.

📘

Download for the following example

To run the following example, first download this pre-trained model.

Python

import coremltools as ct

# Convert a frozen graph from TensorFlow 1 to Core ML
mlmodel = ct.convert("mobilenet_v1_1.0_224/frozen_graph.pb")

The MobileNet model in the previous example already has a defined input shape, so you did not need to provide it. However, in some cases the TensorFlow model does not contain a fully defined input shape. You can pass an input shape that is compatible with the model into the convert() method in order to provide the shape information, as shown in the following example.

📘

Download for the following example

To run the following example, first download this pre-trained model.

Python

import coremltools as ct

# Needs additional shape information
mlmodel = ct.convert("mobilenet_v2_1.0_224_frozen.pb",
                    inputs=[ct.TensorType(shape=(1, 224, 224, 3))])

Convert from PyTorch

You can convert PyTorch models that are either traced or in the TorchScript format. For example, you can convert a model obtained using PyTorch's save and load APIs to Core ML using the same Unified Conversion API as the previous example:

Python

import coremltools as ct
import torch
import torchvision

# Get a pytorch model and save it as a *.pt file
model = torchvision.models.mobilenet_v2()
model.eval()
example_input = torch.rand(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("torchvision_mobilenet_v2.pt")

# Convert the saved PyTorch model to Core ML
mlmodel = ct.convert("torchvision_mobilenet_v2.pt",
                    inputs=[ct.TensorType(shape=(1, 3, 224, 224))])

For more details on tracing and scripting to produce PyTorch models for conversion, see the PyTorch page.

API reference

For the full list of arguments supported by the coremltools.convert API, see the API reference page.

For common scenarios, see the pages on image inputs, classifiers, and flexible input shapes under Conversion Options.

Amazon SageMaker Neo

You can use Amazon SageMaker Neo to automatically convert TensorFlow and PyTorch models to Core ML using the Unified Conversion API. Amazon SageMaker Neo exposes the Unified Conversion API from its service in the AWS cloud. You can access the API using the AWS Command Line Interface (AWS CLI), the Amazon SageMaker console, or the SDK. For more information see the AWS Machine Learning blog.