Importing and Exporting Neural Networks with ONNX

Matteo  Salva rezza/ Maria  Sargsyan/Tuseeta  Banerjee / 
Wolfram Research

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

ABSTRACT

ONNX is an open format to represent deep learning models. It allows easy conversion of models from one framework to other. In this talk we will provide an overview of the current state of the ONNX importer and exporter featured in the Wolfram Language and show how to use their options and properties. With the help of a few models (that are currently in the Neural Net Repository, or will be added soon), we will show how to import and export models from and to the popular deep learning frameworks (Tensorflow-Keras/CoreML/Pytorch). Finally, we will discuss what new features can be expected in the upcoming releases.

Overview: ONNX Import/Export

ONNX (Open Neural Networks eXchange) is a cross-platform specification and file format for machine learning models

It started as a specification for inference in Deep Learning, but has since expanded to training and classic Machine Learning as well

Has widespread support in the industry

Overview: ONNX Import/Export

Import can be used to import an ONNX model into the Wolfram Language and to inspect its metadata:

Import["ExampleData/onnx_example.onnx"]

Import["ExampleData/onnx_example.onnx", "Elements"]

Relevant metadata include:

Import["ExampleData/onnx_example.onnx", "OperatorSetVersion"]

Import["ExampleData/onnx_example.onnx", "IRVersion"]

The full, low-level contents of the ONNX file are also available:

Import["ExampleData/onnx_example.onnx", "RawData"]

It is also possible to only import the arrays:

Import["ExampleData/onnx_example.onnx", "ArrayAssociation"]

Overview: ONNX Import/Export

Nets are exported to onnx via Export:

net=NetInitialize@NetGraph[{LinearLayer[2],Ramp,LinearLayer[3],Tanh,CatenateLayer[]},{NetPort["Input1"]->1->2,NetPort["Input2"]->3->4,{2,4}->5},"Input1"->6,"Input2"->7]

Export["onnx_model.onnx",net]

Roundtrips are usually possible but not guaranteed in all cases:

Import["onnx_model.onnx"]

Some metadata can be set with options:

Export["onnx_model.onnx",net,"GraphName"->"My Graph","ModelVersion"->4,"Description"->"Description here."]

Import["onnx_model.onnx","RawData"]["graph","name"]
Import["onnx_model.onnx","ModelVersion"]
Import["onnx_model.onnx","Description"]

Exporting a net from WNNR into ONNX

Import Forknet into Mathematica

forknet=NetModel["ForkNet Brain Segmentation Net Trained on NAMIC Data"];

Define a brain image and export to a JPEG file:

brainImg=;

Export[FileNameJoin[{$TemporaryDirectory,"brain.jpeg"}],brainImg];

Evaluate the net to get White and Gray Matter:

masks=NetModel["ForkNet Brain Segmentation Net Trained on NAMIC Data"]@brainImg;

Export the net to onnx:

Export[FileNameJoin[{$TemporaryDirectory,"forknet.onnx"}],forknet]

Import the net back to see that the conversion to onnx was correct:

Import["/tmp/forknet.onnx"]

Exporting a net from WNNR into ONNX

Define an external session to run python code:

session=StartExternalSession["Python"];

Define the python code and evaluate it:

code="
import onnx;
import h5py;
from PIL import Image;
import matplotlib.pyplot as plt;
from onnx2keras import onnx_to_keras;
import numpy as np;

model=onnx.load('/tmp/forknet.onnx')
try:
    onnx.checker.check_model(model);
except onnx.checker.ValidationError as e:
    print('The model is invalid: %s' % e);
else:
    print('The model is valid!');
k_model = onnx_to_keras(model, ['Input']);
#k_model.summary()
#Load a brain image, resize to fit the dimensions of the input of the net 
# and convert the input to [0,1] range
brainImg=Image.open('/tmp/brain.jpeg');
keras_input=np.array(brainImg)/255.;
keras_input=keras_input[np.newaxis,np.newaxis,:,:];
keras_output = k_model.predict(keras_input);
print('Prediction is finished');
(wm,gm)=(keras_output[0][0],keras_output[1][0]);
wmImage=Image.fromarray((wm*255).astype('uint8')[0],'L');
gmImage=Image.fromarray((gm*255).astype('uint8')[0],'L');
brainImg.paste(wmImage, (0,0), mask = wmImage);
brainImg.paste(gmImage, (0,0), mask = gmImage);
brainImg.save('/tmp/keras_output_image.jpeg')
";

ExternalEvaluate[session,code]

Visualize the python vs. Mathematica output:

<|"Python Output"-> Import["/tmp/keras_output_image.jpeg"],"Mathematica output"->HighlightImage[brainImg,{Image[masks[["gm",1]]],Image[masks[["wm",1]]]}]|>

Importing a net from Keras into Wolfram

Define the python code and evaluate it to export Efficient-B0 into ONNX:

ExternalEvaluate[session,
"
import tensorflow as tf;
import tf2onnx
model = tf.keras.applications.EfficientNetB0(weights='imagenet',input_shape=(224,224,3));
tf.saved_model.save(model, '/tmp/tmp_model');
onnx_proto,_= tf2onnx.convert.from_keras(model,opset=13);
with open('/tmp/keras_b0.onnx', 'wb') as f:
    f.write(onnx_proto.SerializeToString())
"
];

Import Efficient-B0 into Mathematica:

Import["/tmp/keras_b0.onnx"]

Importing a net from CoreML

(Canvas[(CanvasFrame)])

Obtain a net from CoreML and convert it to ONNX

session=StartExternalSession["Python"]

ExternalEvaluate[session,"import coremltools
import onnxmltools

# Update the input name and path for your CoreML model

input_coreml_model = 'TinyYOLO.mlmodel'

# Change this path to the output name and path for the ONNX model
output_onnx_model = 'TinyYOLO.onnx'

# Load your CoreML model
coreml_model = coremltools.utils.load_spec(input_coreml_model)

# Convert from CoreML into ONNX
onnx_model = onnxmltools.convert_coreml(coreml_model, 'TinyYOLO')

# Save as protobufss
onnxmltools.utils.save_model(onnx_model, output_onnx_model)

import os
print(os.path.getsize(output_onnx_model))"]

This is very similar to 
https://gallery.azure.ai/Model/Tiny-YOLOv2

Importing a net from CoreML

(Canvas[(CanvasFrame)])

Import the WL and complete missing sections

tinyyolo=
Import["TinyYOLO.onnx"]

tinyyolosurgery = NetTake[tinyyolo,{{"convolution"},{"convolution8"}}]

convNet[]:=tinyyolosurgery

regionLayerNet[gridSize:{w_,h_},nAnchors_,nCoords_,nClasses_]:=NetGraph[
{
"Reshape"->ReshapeLayer[{nAnchors,nCoords+1+nClasses,w,h}],
"Transpose"->TransposeLayer[],
"Centers"->{PartLayer[1;;2],ElementwiseLayer["Sigmoid"]},
"Sizes"->PartLayer[3;;4],
"Coordinates"->CatenateLayer[],
"Confidence"->{PartLayer[nCoords+1],ElementwiseLayer["Sigmoid"],FlattenLayer[]},
"ClassScore"->{PartLayer[nCoords+2;;All],SoftmaxLayer[1],FlattenLayer[-2],TransposeLayer[]}
},
{"Reshape"->"Transpose",
"Transpose"->"Centers"->"Coordinates",
"Transpose"->"Sizes"->"Coordinates",
"Coordinates"->NetPort["Coordinates"],
"Transpose"->"Confidence"->NetPort["Objectness"],
"Transpose"->"ClassScore"->NetPort["ClassProb"]},
"Input"->{nAnchors (nCoords+1+nClasses),w,h}
];

coord2box[gridSize:{w_,h_},anchors_,nCoords_,nClasses_]:=With[
{
grid=Transpose[Table[{i,j},{j,0,w-1},{i,0,h-1}],{2,3,1}],
n=Length[anchors]
},
With[
{
gridLayer=NetArrayLayer["Array"->ReplicateLayer[n,2][grid]],
anchorLayer=NetArrayLayer["Array"->ReplicateLayer[{w,h},3][Transpose[anchors]]],
scaleLayer=NetArrayLayer["Array"->ReplicateLayer[{n,w,h},2][{w,h,w,h}]]
},
NetGraph[
{"Centers"->PartLayer[1;;2,"Input"->{nCoords,n,w,h}],
"ScaledCenters"->TotalLayer[],
"Sizes"->{PartLayer[3;;4],ElementwiseLayer[Exp]},
"ScaledSizes"->ThreadingLayer[Times],
"ScaledCoord"->CatenateLayer[],
"NormalizedCoord"->ThreadingLayer[Divide],
"FlattenedCoord"->FlattenLayer[-2],
"CoordList"->TransposeLayer[],
"Grid"->gridLayer,
"Scale"->scaleLayer,
"Anchors"->anchorLayer},
{
NetPort["Input"]->"Centers"->"ScaledCenters"->"ScaledCoord",
"Grid"->"ScaledCenters",
NetPort["Input"]->"Sizes"->"ScaledSizes"->"ScaledCoord",
"Anchors"->"ScaledSizes",
{"ScaledCoord","Scale"}->"NormalizedCoord"->"FlattenedCoord"->"CoordList"
},
"Input"->{nCoords,n,w,h}
]
]
];

scaleImage:=Function[img,
ImageData[ImageCrop[ColorConvert[
ImageResize[RemoveAlphaChannel[img],{416},Resampling->"Linear"],
"RGB"],{416,416},Padding->.5],Interleaving->False]]

anchors=(*{{0.57273,0.677385},{1.87446,2.06253},{3.33843,5.47434},{7.88282,3.52778},{9.77052,9.16828}};*)
anchors={{1.08,1.19},{3.42,4.41},{6.63,11.38},{9.42,5.11},{16.62,10.52}};

w=h=13;
nclasses=20;
ncoord=4;

Final Net

tinyyoloV2Net=NetGraph[
{
"Conv"->convNet[],
"Region"->regionLayerNet[{w,h},Length@anchors,ncoord,nclasses],
"BoxTransformation"->coord2box[{w,h},anchors,ncoord,nclasses]
},
{
"Conv"->"Region",
NetPort["Region","Coordinates"]->"BoxTransformation"->NetPort["Boxes"]
},
"Input"->NetEncoder[{"Function",scaleImage,{3,416,416}}]
]

Importing a net from CoreML

Evaluate model

Write an evaluation function to scale the result to the input image size and suppress the least probable detections:

nonMaxSuppression[overlapThreshold_][detection_] := 
Module[{boxes, confidence},
	Fold[
		{list, new} |-> If[NoneTrue[list[[All, 1]], IoU[#, new[[1]]] > overlapThreshold &], Append[list, new], list],
		Sequence @@ TakeDrop[Reverse@SortBy[detection, Last], 1]
	]
]
ClearAll[IoU]
IoU := IoU =  
With[{c = Compile[
	{{box1, _Real, 2}, {box2, _Real, 2}},
	Module[{area1, area2, x1, y1, x2, y2, w, h, int},
	
		area1 = (box1[[2, 1]] - box1[[1, 1]])(box1[[2, 2]] - box1[[1, 2]]);
		area2 = (box2[[2, 1]] - box2[[1, 1]])(box2[[2, 2]] - box2[[1, 2]]);
		
		x1 = Max[box1[[1,1]], box2[[1,1]]];
		y1 = Max[box1[[1,2]], box2[[1,2]]];
		x2 = Min[box1[[2,1]], box2[[2,1]]];
		y2 = Min[box1[[2,2]], box2[[2,2]]];	
	
		w = Max[0., x2 - x1];
		h = Max[0., y2 - y1];
	
		int = w * h;
		
		int / (area1 + area2 - int)
	],
	RuntimeAttributes -> {Listable},
	Parallelization -> True,
	RuntimeOptions -> "Speed"
]},
	c @@ Replace[{##}, Rectangle -> List, Infinity, Heads -> True]&
]
netevaluate[img_, detectionThreshold_ : .2, overlapThreshold_ : .2] :=
Module[
  	{w, h, scale, coords, obj, classes, boxes, padding, classProb, bestClass, probable, finals},
  
  	{obj, classes, coords} = Values@tinyyoloV2Net[img];
  
  	(* transform coordinates into rectangular boxes *)
  	{w, h} = ImageDimensions[img];
  	scale = Max[{w, h}] / 416;
  	padding = 416 (1 - {w, h} / Max[{w, h}]) / 2;
  	boxes = Apply[
    		Rectangle[
      			scale {#1 - #3/2 - padding[[1]], 416 - #2 - #4/2 - padding[[2]]},
      			scale {#1 + #3/2 - padding[[1]], 416 - #2 + #4/2 - padding[[2]]}
      		] &,
    		416 coords,
    	1];
  	
  	(* each class probability is rescaled with the box objectivness *)
  	classProb = classes * obj;
  	(* filter by probability*)
  	(* very small probability are thresholded *)
  	probable = Position[Max /@ classProb , p_ /; p - 10^-2 > detectionThreshold];	
  	If[Length[probable] == 0, Return[{}]];
  	
  	(* gather the boxes of the same class and perform non-max suppression *)
  	bestClass = Last@*Ordering /@classProb;
  	finals = Join @@ Values @ GroupBy[
      		MapThread[
       			{#1, #2, #3[[#2]]} &,
       			{Extract[boxes, probable], Extract[bestClass, probable], Extract[classProb, probable]}
       		],
      		#[[2]]&, nonMaxSuppression[overlapThreshold]
      	]
  ]

labels={"aeroplane","bicycle","bird","boat","bottle","bus","car","cat","chair","cow","diningtable","dog","horse","motorbike","person","pottedplant","sheep","sofa","train","tvmonitor"};

testImage=;

detection=netevaluate[testImage]

classes=DeleteDuplicates@detection[[All,2]]

labels[[classes]]

HighlightImage[testImage,MapThread[{White,Inset[Style[labels[[#2]],Black,FontSize->Scaled[1/12],Background->GrayLevel[1,.6]],Last[#1],{Right,Top}],#1}&,Transpose@detection]]

Exporting a net from WNNR to CoreML

(Canvas[(CanvasFrame)])

Obtain the net from WNNR and export the necessary files to ONNX

pretrainedNet=NetModel["MobileNet V2 Trained on ImageNet Competition Data"]

Export["MobileNetV2.onnx",pretrainedNet];

classLabels=EntityValue[NetExtract[NetModel["MobileNet V2 Trained on ImageNet Competition Data"],"Output"][["Labels"]],"Name"];

Export["ImageNetClassLabels.txt",StringRiffle[classLabels,"\n"]];

Exporting a net from WNNR to CoreML

(Canvas[(CanvasFrame)])

Convert the ONNX model to .mlmodel

ExternalEvaluate[session,"import coremltools as ct
# Specify path to ONNX model
onnx_model_path = 'MobileNetV2.onnx'
# Convert from ONNX to Core ML
mlmodel =  ct.converters.onnx.convert(onnx_model_path,
                                      image_input_names = ['Input'],
                                      mode = 'classifier',
                                      preprocessing_args ={'image_scale': 1.0/255.0,'red_bias': -0.485,'green_bias': -0.456,'blue_bias': -0.406},
									 class_labels = 'ImageNetClassLabels.txt',
                                      minimum_ios_deployment_target = '13')
                                    
# Save the mlmodel
mlmodel.save('MobileNetV2.mlmodel')"]

Exporting a net from WNNR to CoreML

Evaluate the model

ExternalEvaluate[session,"import coremltools

model = coremltools.models.MLModel('MobileNetV2.mlmodel')

import numpy as np
import PIL

img = PIL.Image.open('cat.jpg')

spec = model._spec
img_width = spec.description.input[0].type.imageType.width
img_height = spec.description.input[0].type.imageType.height
img = img.resize((img_width, img_height), PIL.Image.BILINEAR)

y = model.predict({'Input': img}, usesCPUOnly=True)

def printTop3(resultsDict):

    # Put probabilities and labels into their own lists.

    probs = np.array(list(resultsDict.values()))
    labels = list(resultsDict.keys())

    # Find the indices of the 3 classes with the highest probabilities.

    top3Probs = probs.argsort()[-3:][::-1]

    # Find the corresponding labels and probabilities.

    top3Results = map(lambda x: (labels[x], probs[x]), top3Probs)

    # Print them from high to low.

    for label, prob in top3Results:
        print('%.5f %s' % (prob, label))

printTop3(y['Output'])"]

img=Import@"cat.jpg";

NetModel["MobileNet V2 Trained on ImageNet Competition Data"][img,{"TopProbabilities"-> 3}]

(Entity["peacock", Entity["Concept","Peacock::kd2y4"], "Entity[\"Concept\", \"Peacock::kd2y4\"]", "concept"])["Properties"]

(Entity["peacock", Entity["Concept","Peacock::kd2y4"], "Entity[\"Concept\", \"Peacock::kd2y4\"]", "concept"])["Definition"]

Although the probability are different (owing to absence of variance corrections, the results match!)

Please note this section is based on:
https://blog.wolfram.com/2021/01/07/deploy-a-neural-network-to-your-ios-device-using-the-wolfram-language/

Importing a net from Pytorch

Convert a net from Pytorch model to ONNX

ExternalEvaluate[session,"import torch
import torchvision
import torch.onnx
import torch.nn as nn
# This script loads a pretrained video resnet model and then exports its sub-blocks to onnx. 

model = torchvision.models.video.r2plus1d_18(pretrained=True)
head =\"./r2plus1d_18\"

model_children = list(model.children())
fe = nn.Sequential(*model_children[:-2])
fe.eval()
torch.onnx._export(fe,            
                           (torch.rand(1,3, 5, 64,64)),                   
                           head+\"_fe.onnx\",
                           #training=torch.onnx.TrainingMode.PRESERVE,
                           #opset_version=13,
                           do_constant_folding=False
                           ) 

torch.onnx._export(model_children[-1].eval(),             
                           torch.rand(1,512),                    
                          head+\"_fc.onnx\",
                          )"]

w=h=13;
nclasses=20;
ncoord=4;

Importing a net from Pytorch

Import the net from ONNX, do necessary surgery

backbone=Import@"r2plus1d_18_fe.onnx";

linear=Import@"r2plus1d_18_fc.onnx";

net = NetChain[Join[
		{
			"transpose" -> TransposeLayer[],
			"stem" -> backbone
		},
		{
			"prediction" -> NetChain[{
				"final_Relu" -> ElementwiseLayer[Ramp], 
				"global_pool" -> AggregationLayer[Mean, {2 ;; All}], 
				"fc" -> linear
			}],
			"softmax" -> SoftmaxLayer[]
		}
	]]

fnet=NetReplacePart[
		net,
		{
			"Input" -> NetEncoder[{
				"VideoFrames", {112, 112}, "TargetLength" -> 16, FrameRate -> Inherited, "ColorSpace" -> "RGB", 
				"VarianceImage" -> {0.22803, 0.22145, 0.216989}, "MeanImage" -> {0.43216, 0.394666, 0.37645}
			}],
			"Output" -> NetDecoder[{"Class", Import["Kinetics400.m"]}]
		}
	]

Importing a net from Pytorch

Evaluate Model

video=ResourceData["Sample Video: Barbecuing"]

fnet[video]

Exporting a net to Pytorch

(Canvas[(CanvasFrame)])

Obtain the net from WNNR and export the necessary files to ONNX

pretrainedNet=NetModel["ShuffleNet-V2 Trained on ImageNet Competition Data"]

Export["shufflenetV2.onnx",pretrainedNet];

classLabels=EntityValue[NetExtract[NetModel["ShuffleNet-V2 Trained on ImageNet Competition Data"],"Output"][["Labels"]],"Name"];

Export["ImageNetClassLabels_ShuffleNetV2.txt",StringRiffle[classLabels,"\n"]];

Exporting a net to Pytorch

(Canvas[(CanvasFrame)])

Convert the ONNX model to Pytorch and evaluate

ExternalEvaluate[session,"import onnx
from onnx2pytorch import ConvertModel
import torch
from torchvision import transforms
from PIL import Image
import numpy as np

#Load ONNX Model
onnx_model = onnx.load('shufflenetV2.onnx')

# Convert from ONNX to Pytorch
pytorch_model = ConvertModel(onnx_model)
pytorch_model.eval()
                                   
# Input and Pre-Process an image
input_image = Image.open('cat.jpg')
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_1 = input_tensor.unsqueeze(0)

# Evaluate model 
output_1 = pytorch_model(input_1)
labels = open('ImageNetClassLabels_ShuffleNet.txt').read().splitlines()

#Compute top 3 probability
def printTop3(resultsDict):

    # Put probabilities and labels into their own lists.

    probs = output_1.detach().numpy()

    # Find the indices of the 3 classes with the highest probabilities.

    top3Probs = probs.argsort(axis=None)[-3:][::-1]

    # Find the corresponding labels and probabilities.

    top3Probs = probs.argsort(axis=None)[-3:][::-1]
    top3Results = map(lambda x: labels[x], list(top3Probs))
    print(list(top3Results))

printTop3(output_1)"]

img=Import@"cat.jpg";

NetModel["MobileNet V2 Trained on ImageNet Competition Data"][img,{"TopProbabilities"-> 3}]

Future

Keep expanding layer support for Import and Export (expect 100% layer support for Export in WL 13.1)

Improve challenging cases for the importer

Some of our layers have static parameters (e.g. weights in a ConvolutionLayer or output shape in a ReshapeLayer) that ONNX treats as extra layer inputs

NetInitialize@ConvolutionLayer[2,3,"Input"->{3,10,10}]

ReshapeLayer[{2,5,3,4},"Input"->{8,1,15}]

In case these extra layer inputs depend on the model’s inputs (i.e. if they change when the model input changes), we fail to import the model

ONNXUtilities`ONNXExport["unsupported_model.onnx",]

Import["unsupported_model.onnx"]

Allow exporting to multiple operator set versions

We usually update the exporter to the latest operator set version at every new WL release, but since other frameworks might be slower to adopt the new version (or stop supporting new versions entirely) the ability to export to past operator set versions is important:

Export["model.onnx",net,"OperatorSetVersion"->10]