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.
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#
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
# 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
# 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 examples of converting PyTorch models, see the following: