PyTorch Conversion

Starting with coremltools 4.0, you can convert your model trained in PyTorch to the Core ML format directly, without requiring an explicit step to save the PyTorch model in ONNX format. Converting the model directly is recommended.

📘

Minimum deployment target

The Unified Conversion API produces Core ML models for iOS 13, macOS 10.15, watchOS 6, tvOS 13 or newer deployment targets.

If your primary deployment target is iOS 12 or earlier, you can find limited conversion support for PyTorch models via the onnx-coreml package.

To first create a representation of a model from PyTorch code, use TorchScript. You can get the TorchScript version of your model with combination of torch.jit.trace calls or directly annotating your model with torch.jit.script calls.

📘

For more information

• To learn how TorchScript works, see the Introduction to TorchScript.
• Use PyTorch's JIT tracer to generate TorchScript for your PyTorch model. To learn how, see Model Tracing.
• If model tracing is inappropriate for your PyTorch model, annotate the model to generate TorchScript directly. To learn how, see Model Scripting.
• For an example of converting a PyTorch segmentation model, see Convert a PyTorch Segmentation Model.

Steps for the conversion example

The following example shows how to convert into Core ML a MobileNetV2 model trained using PyTorch. The example is similar to the one provided in the Quickstart Example, in which you convert the TensorFlow version of the model.

In this example you do the following:

1. Load the model and set it to evaluation mode.
2. Generate TorchScript using the torch.jit.trace command provided in PyTorch.
3. Convert the TorchScript object to Core ML using the Unified Conversion API.

Once you have converted the model, you can follow the steps from the Quickstart Example to evaluate and test your model.

Load the MobileNetV2 model

The example uses a pre-trained version of the MobileNetV2 model from the Torchvision library. Follow these steps:

1. Load the pre-trained version of MobileNetV2:

Python

import torch
import torchvision

# Load a pre-trained version of MobileNetV2
torch_model = torchvision.models.mobilenet_v2(pretrained=True)

2. Set the model to evaluation mode:

Python

# Set the model in evaluation mode
torch_model.eval()

📘

Set the model to evaluation mode

To ensure that operations such as dropout are disabled, it's important to set the model to evaluation mode (not training mode) before tracing. This setting also results in a more optimized version of the model for conversion.

Generate TorchScript by tracing the model

The process of tracing takes an example input and traces its flow through the model. You can trace the model by creating an example image input. The rank and shape of the tensor depends on your model's use case. If your model expects a fixed-size input, use that size for the example image. If your model can accept a variety of input sizes, the best approach is to keep the example image small in order to reduce the time it takes to trace. In all cases, the rank of the tensor must be fixed.

Follow these steps:

1. Trace the model with random data.

Python

# Trace with random data
example_input = torch.rand(1, 3, 224, 224) # after test, will get 'size mismatch' error message with size 256x256
traced_model = torch.jit.trace(torch_model, example_input)

2. [Optional] Download the class labels from a separate file:

Python

# Download class labels (from a separate file)
import urllib
label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'
class_labels = urllib.request.urlopen(label_url).read().decode("utf-8").splitlines()

class_labels = class_labels[1:] # remove the first class which is background
assert len(class_labels) == 1000

Convert the TorchScript object to Core ML

Convert the model to Core ML using the Unified Conversion API:

Python

import coremltools as ct
# Convert to Core ML using the Unified Conversion API
model = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input_1", shape=example_input.shape)], #name "input_1" is used in 'quickstart'
    classifier_config = ct.ClassifierConfig(class_labels) # provide only if step 2 was performed
)

📘

Images as inputs

By default, the coremltools converter generates a Core ML Model with inputs of type MLMultiArray. By providing additional input arguments as shown in the previous example, you can use either TensorType or ImageType. To learn how to work with images as inputs and achieve better performance and more convenience, see Image Inputs.

With the converted model in Core ML memory, you can save the model into the Core ML format:

Python

# Save model
model.save("MobileNetV2.mlmodel")

Once you have converted and saved the model, you can follow the steps in the Quickstart Example to:

• Add metadata to the model.
• Make predictions using the model.
• Use the model with Xcode.