You can convert a 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. (This feature was introduced in Core ML Tools version 4.0.)

📘

Minimum Deployment Target

The Core ML Tools 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.

Generate a TorchScript Version

TorchScript is an intermediate representation of a PyTorch model. To generate a TorchScript representation from PyTorch code, use PyTorch's JIT tracer (torch.jit.trace) to trace the model, as shown in the following example:

import torch
import torchvision

# Load a pre-trained version of MobileNetV2
torch_model = torchvision.models.mobilenet_v2(pretrained=True)
# Set the model in evaluation mode.
torch_model.eval()

# Trace the model with random data.
example_input = torch.rand(1, 3, 224, 224) 
traced_model = torch.jit.trace(torch_model, example_input)
out = traced_model(example_input)

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, as shown in the above code using random data. To understand the reasons for tracing and how to trace a PyTorch model, see Model Tracing.

📘

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.

If your model uses a data-dependent control flow, such as a loop or conditional, the traced model won't generalize to other inputs. In such cases you can experiment with applying PyTorch's JIT script (torch.jit.script) to your model as described in Model Scripting. You can also use a combination of tracing and scripting.

Convert to Core ML

Convert the traced or scripted model to Core ML using the Unified Conversion API convert() method. In the inputs parameter, you can use either TensorType or ImageType as the input type.

The following example uses TensorType and converts the PyTorch traced model to a Core ML program model. For more about image input conversions, see Image Inputs.

import coremltools as ct

# Using image_input in the inputs parameter:
# Convert to Core ML program using the Unified Conversion API.
model = ct.convert(
    traced_model,
    convert_to="mlprogram",
    inputs=[ct.TensorType(shape=example_input.shape)]
 )

With the converted ML model in memory, you can save it as a Core ML model package:

# Save the converted model.
model.save("newmodel.mlpackage")

As an alternative, you can convert the model to a neural network by eliminating the convert_to parameter:

# Using image_input in the inputs parameter:
# Convert to Core ML neural network using the Unified Conversion API.
model = ct.convert(
    traced_model,
    inputs=[ct.TensorType(shape=example_input.shape)]
 )

With the converted neural network in memory, you can save it as an mlmodel file:

# Save the converted model.
model.save("newmodel.mlmodel")

📘

For More Information

• To learn how TorchScript works, see the Introduction to TorchScript.
• To learn how to get better performance and more convenience when using images, see Image Input and Output.
• For the details of how to preprocess image input for models in PyTorch's torchvision library, see Preprocessing for Torch.
• For examples of converting PyTorch models, see the following:
  • Converting a Natural Language Processing Model
  • Converting a torchvision Model from PyTorch
  • Converting a PyTorch Segmentation Model