Skip to content
A simple guide that explains the steps from training a simple PyTorch image classifier to converting the generated neural network into a CoreML model ready for production
Python
Branch: master
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
LICENSE
README.md
model.py
my_network.mlmodel
my_network.onnx
my_network.pth
step1.md
step1.py
step2.md
step2.py
step3.md
step3.png
step3.py
step4.md
step4.png
step4.py
step5.md
step5.py

README.md

Motivation

I present here a simple guide that explains the steps needed from training a simple PyTorch image classifier to converting the trained neural network into a CoreML model ready for production. I've spent days crawling Internet blogs, forums and official documentations to gather the little knowledge presented in these pages. The true motivation of this repo is to prevent me to forget everything I know about this particular subject. If this guide helps someone else to move forward into her/his research, that's a plus to me. Please read the disclaimer.

The guide

The problem I faced was pretty simple. I wanted to know how to train an artificial neural network in PyTorch and how to convert this network into a CoreML model usable in an iOS application. Simple right?

I present bellow my solution divided into several steps:

  • Step 1: Train a model using PyTorch and save it
  • Step 2: Load the model and test it (for verification purpose)
  • Step 3: Convert the PyTorch model into a ONNX model
  • Step 4: Convert the ONNX model into a CoreML model
  • Step 5: Test the CoreML model

Virtual environment

For this work, I've been using Conda through Anaconda for (1) creating a virtual environment and (2) installing most of the used Python packages. Please read the official documentations for more information.

Disclaimer

As presented in the motivation section, the target audience of this guide is me. I am of course happy if it helps other coders around the world. I do not certify the code presented is the best or even the correct way to using PyTorch, ONNX, coremltools, etc. The code is probably uncomplete and might even contain serious bugs. In addition, this code will probably ceased to work with the newer versions of the different libraries and the evolution of PyTorch itself. Use this code at your own risk.

The code presented here has been written in late 2019 on a MacBook Pro using:

  • torch 1.1.0.post2
  • pytorch 1.3.0
  • torchvision 0.3.0
  • onnx 1.5.0
  • onnx-coreml 1.0
  • coremltools 3.0
  • numpy 1.16.4
  • pillow 6.0.0

If you see something wrong, please let me know and I'll be happy to make modifications.

Acknowledgment

I've been reading intensively PyTorch tutorials to educate myself. Don't be surprised if you find some similarities in the code.

You can’t perform that action at this time.