|
| 1 | +### Converting PyTorch model to CoreML model (Short version) |
| 2 | +Converting PyTorch model into CoreML model is a two step process |
| 3 | +1. Convert PyTorch model into ONNX model |
| 4 | + - PyTorch model can be converted into ONNX model using `torch.onnx.export` |
| 5 | + - Reference: https://pytorch.org/docs/stable/onnx.html#id2 |
| 6 | + - Tools required: [PyTorch](https://pytorch.org/get-started/locally/) |
| 7 | +2. Converting ONNX model into CoreML model |
| 8 | + - We will be using model converted in step 1 using `onnx_coreml.convert` |
| 9 | + - Tools required: [onnx-coreml](https://pypi.org/project/onnx-coreml/) |
| 10 | + |
| 11 | + |
| 12 | +### Converting PyTorch model to CoreML model (Long version) |
| 13 | +1. PyTorch to ONNX conversion |
| 14 | + - In this example, we are creating a small test model |
| 15 | + ``` |
| 16 | + import torch |
| 17 | + import torch.nn as nn |
| 18 | + import torch.nn.functional as F |
| 19 | +
|
| 20 | + # Step 0 - (a) Define ML Model |
| 21 | + class small_model(nn.Module): |
| 22 | + def __init__(self): |
| 23 | + super(small_model, self).__init__() |
| 24 | + self.fc1 = nn.Linear(768, 256) |
| 25 | + self.fc2 = nn.Linear(256, 10) |
| 26 | +
|
| 27 | + def forward(self, x): |
| 28 | + y = F.relu(self.fc1(x)) |
| 29 | + y = F.softmax(self.fc2(y)) |
| 30 | + return y |
| 31 | + ``` |
| 32 | + - Load model |
| 33 | + ``` |
| 34 | + # Step 0 - (b) Create model or Load from distk |
| 35 | + model = small_model() |
| 36 | + dummy_input = torch.randn(768) |
| 37 | + ``` |
| 38 | + - Convert From PyTorch to ONNX |
| 39 | + ``` |
| 40 | + # Step 1 - PyTorch to ONNX model |
| 41 | + torch.onnx.export(model, dummy_input, './small_model.onnx') |
| 42 | + ``` |
| 43 | +2. ONNX to CoreML |
| 44 | + ``` |
| 45 | + # Step 2 - ONNX to CoreML model |
| 46 | + mlmodel = convert(model='./small_model.onnx', target_ios='13') |
| 47 | + # Save converted CoreML model |
| 48 | + mlmodel.save('small_model.mlmodel') |
| 49 | + ``` |
| 50 | +
|
| 51 | +### What about frameworks other than PyTorch? |
| 52 | + - Step 1 can be replaced by respective framework to ONNX converter, |
| 53 | + - Once, you have a ONNX model, you can use onnx_coreml.convert to get CoreML model |
| 54 | + which can be deployed on device. |
| 55 | + |
| 56 | +### More examples demonstrating prediction and validating converted model's correctness |
| 57 | + - BERT: https://github.com/onnx/onnx-coreml/blob/master/examples/BERT.ipynb |
| 58 | + - All: https://github.com/onnx/onnx-coreml/blob/master/examples/ |
0 commit comments