convert_onnx_coreml.py

import torch
import torch.nn as nn
import numpy as np
from onnx_coreml import convert
from torch.autograd import Variable
import torch.onnx
import torchvision
import onnx
 
model_name = "dog_vs_cat_image.onnx"

# Pytorch model --> onnx model & check if it is convertible
dummy_input = Variable(torch.randn(3, sz, sz)).cuda() # for iOS app, we predict only 1 image at a time, we don't use batch
torch.onnx.export(model, dummy_input, model_name, \
                  input_names=['image'], output_names=['dogorcat'], verbose=True)
onnx_model = onnx.load(model_name)

# onnx model --> Apple Core ML
mlmodel = convert(onnx.load(model_name), image_input_names = ['image'], \
                  mode='classifier', class_labels="labels.txt")
mlmodel.author = 'jalola'
mlmodel.license = 'MIT'
mlmodel.short_description = 'This model takes a picture of a dog or cat and predicts its a cat or a dog'
mlmodel.input_description['image'] = 'Image of a dog or cat or something else'
mlmodel.output_description['dogorcat'] = 'Confidence and label of predicted dog or cat'
mlmodel.output_description['classLabel'] = 'Label of predicted dog or cat'
mlmodel.save(f'{model_name}.mlmodel')