PyTorch Conversion tips when using the CoreML Framework and vision framework

Photo by Willian Justen de Vasconcellos on Unsplash

CoreML is a great framework apple released in 2017, compared with Google it’s nearly one year earlier. But as we all know, google publish TensorFlow, so if you want to apply the machine learning into your application, TensorFlow lite is your choice.

PyTorch is a popular framework in DL. Compared with TensorFlow it’s easier to start your DL project because the learning curve is very smooth.

But if you want to use CoreML in your iOS application, you must convert it to the model (coreml). But unfortunately, it cannot convert to CoreML model directly. There are two steps to go:

1. Convert PyTorch to onnx
2. Convert onnx to CoreML

First, instantiate the PyTorch model and load its weights from a checkpoint.

Unlike models from other frameworks, PyTorch models usually don’t include the definition of the model’s architecture, only the learned weights. It means you need to have access to the original Python source code that defines the model.

If the model contains CUDA tensors- which can only be loaded if your computer has an NVIDIA GPU, and most Macs don’t — then load the checkpoints as follows.

checkpoint = torch.load("my_model.pt",map_location = "cpu")
model.load_state_dict(checkpoint)

Now that model is loaded, you can convert it to ONNX format

dummy_input = torch.randn(1,3,224,224)
torch.onnx.export(model, dummy_input,"myModel.onnx",verbose=True)

To do the conversion, PyTorch will perform a forward pass through the model.

That’s why you need to provide the dummy_input tensor.

If your model’s forward() takes multiple inputs, dummy_input should be a tuple of tensors.

Note that in PyTorch the shape is (Batch size, channels, height, and width)

By default, onnx.export() assigns numeric identifiers to the model’s inputs and outputs(and layer names). That’s not very nice, which is why you can also provide human-readable names

torch.onnx.export(model, dummy_input,
"myModel.onnx",verbose=True,
               input_names=["image"],
output_names=["prediction"])

Use verbose=True to print out the layers from the converted model. This is useful for confirming that the model conversion was successful. Plus it shows the size of the intermediate tensors, which is always good to check.

If at this point you get a nasty looking error message, it usually means the PyTorch model performs an operation that is not supported by ONNX converter. Usually, the error is “RuntimeError:ONNX export failed”.

ONNX doesn’t support all operations from PyTorch, and conversely, torch.onnx.export may not support all operations that are possible in ONNX.

The PyTorch documentation lists the operations that torch.onnx.export can handle.

If your model has layers that ONNX doesn’t support, and they are at the very end of the model, you can convert up to the last supported layer and implement the rest of the model yourself as custom layers in CoreML.

When you do custom layers in CoreML, you can add a layer at the beginning of the model/at the end of the model/in the middle of the model, you can delete the layers, you can add softmax layer in your model.

If the unsupported operations are somewhere in the middle, you may need to split up the model into two or more smaller models and convert those separately, then glue them back together after converting to Core ML. Unfortunately, there is currently no way to add placeholders for the unsupported operations in the ONNX file.

Exporting a partial PyTorch model is very easy to do. Since you have the model source code anyway, just comment out any of the unsupported layers in the forward() method, like that:

class MyModel(torch.nn.Module):
   ...
    def forward(self, x):
      x = self.layer1(x)
      x = self.layer2(x)
      return x

If all goes well, you should have an ONNX model.

Import: PyTorch models usually don’t include operations for preprocessing the input, and so there are not added to the ONNX model. If your model does image preprocessing, don’t forget to manually set these options when you convert the onnx Model to Core ML.

from onnx_coreml import convert 
coreml_model = convert("myModel.onnx")
coreml_model.save("myModel.mlmodel")

For models that work on images, it’s common to normalize the pixel values. You’ll need to provide the preprocessing_args argument, which is a dictionary containing one or more of the following values:

• image_scale:
• red_bias/blue_bias/green_bias
• gray_bias
• is_bgr

For models that have image outputs, you can also provide a similar dictionary named deprocessing_args. These “de-processing” options are applied to the outputs listed in image_output_names. This is useful for converting the pixels in the predicted images to the range 0–255.