PyTorch Conversion Workflow

PyTorch Conversion Workflow#

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.

To export a model from PyTorch to Core ML, there are 2 steps:

  1. Capture the PyTorch model graph from the original torch.nn.Module, via torch.jit.trace or torch.export.export.

  2. Convert the PyTorch model graph to Core ML, via the Core ML Tools Unified Conversion API.

The conversion from a graph captured via torch.jit.trace has been supported for many versions of Core ML Tools, hence it is the stable and the more performant path. For now, this is the recommended way for converting PyTorch models to Core ML.

The conversion from torch.export graph has been newly added to Core ML Tools 8.0. It is currently in beta state, in line with the export API status in PyTorch. As of Core ML Tools 8.0, representative models such as MobileBert, ResNet, ViT, MobileNet, DeepLab, OpenELM can be converted, and the total PyTorch op translation test coverage is roughly ~70%. You can start trying the torch.export path on your models that are working with torch.jit.trace already, so as to gradually move them to the export path as PyTorch also moves its support and development to that path over a period of time. In case you hit issues (e.g. models converted via export path are slower than the ones converted from jit.trace path), please report them on Github.

Now let us take a closer look at how to convert from PyTorch to Core ML through an example.

Obtain the Original PyTorch Model#

As an illustration, here we define a toy PyTorch model

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()

Set the Model to Evaluation Mode

To ensure that training time operations such as dropout, batch norm with moving average etc. are disabled, it is important to set the model to evaluation mode before tracing / exporting. This setting also results in a more optimized version of the model for deployment.

Capture the PyTorch Graph#

TorchScript#

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:

# Trace the model with random data.
example_input = torch.rand(1, 3, 224, 224) 
traced_model = torch.jit.trace(torch_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.

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 use JIT script (torch.jit.script) to capture the graph accurately. Core ML Tools offers limited support for models that are obtained from torch.jit.script API, so it may or may not work for your model. See Model Scripting to learn more.

ExportedProgram#

ExportedProgram is a new intermediate representation of a PyTorch model introduced in PyTorch 2. To generate an ExportedProgram representation from PyTorch code, use PyTorch’s exporter (torch.export.export) to export the model, as shown in the following example:

# Export the model with random data.
example_inputs = (torch.rand(1, 3, 224, 224),)
exported_program = torch.export.export(torch_model, example_inputs)

The process of exporting 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.

Convert to Core ML#

Convert the traced / exported 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 for the traced model.

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

import coremltools as ct

# Convert to Core ML program using the Unified Conversion API.
model_from_trace = ct.convert(
    traced_model,
    inputs=[ct.TensorType(shape=example_input.shape)],
)
model_from_export = ct.convert(exported_program)

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

# Save the converted model.
model_from_trace.save("newmodel_from_trace.mlpackage")
model_from_export.save("newmodel_from_export.mlpackage")

For More Information