Exchanging Neural Networks Using ONNX
업데이트:
Open Neural Network Exchange: ONNX
Problems and Solution
Tensorflow, PyTorch, Caffe, Keras, CoreML, there exist a bunch of different frameworks for artificial intelligence. ONNX enables to unlock the doors between different frameworks, prevent the researchers from being locked in one ecosystem. Their goal is to make it possible for users to use the right combination of tools for their project.
Open Neural Network Exchange (ONNX) is a powerful tool used as rerpresenting the deep learning model. Its a platform for different frameworks to have a compatible neural network. For example, lets assume a model is created based on Tensorflow framework. If we export the model using ONNX, we can import the model on PyTorch using ONNX.
There are some frameworks that support ONNX. The official documents are provided.
Strengths and Limitations
- Strengths
- Framework Interpretability: ONNX makes the neural network compatible between different frameworks.
- Shared Optimization: If there exists a shared format between frameworks, the optimization on hardware for neural network can be more easily done. The optimization could be implemented based on the ONNX structure.
- Limitation
- Incompleteness of JIT (Just-In-Time) compiler makes the support of ONNX partial.
- Difference in backend implementation cause the problem in model performance.
Using ONNX from PyTorch
PyTorch provides the module to use ONNX from framework itself. The below example trains pretrained ResNet18 with CIFAR10 dataset, and save model in two ways. One is saved using general torch.save and the other is saved using general torch.onnx.export.
The code for training and comparing is avaliable in https://github.com/sanglee325/onnx-pytorch.
Save Models
-
Save model in general PyTorch weight model.
CKPT_PATH = './ckpt/resnet18.pth' torch.save(model.state_dict(), CKPT_PATH) -
Save model as ONNX in PyTorch.
# to save in ONNX format, dummy input size is required args = torch.randn(1, 3, 32, 32, device=device) ONNX_PATH = './ckpt/resnet18.onnx' torch.onnx.export(model, args, ONNX_PATH, verbose=False)
Compare ONNX and PyTorch
-
Run
compare.pyfrom the https://github.com/sanglee325/onnx-pytorch. If the ONNX and PyTorch weight matches, the result will be shown as below.psl@icsl-53:~/onnx-pytorch$ python compare.py fc.weight: no difference.
Movie Stream Scenario
When we were doing the group project, movie recommendation service, there were many codes on GitHub as an open source. However, the framework of each code was very different making it hard for us to try test each models, that uses different framework. The environment was set for PyTorch but the one with better performance were using Keras. So we had to give up using the better model for the task.
Usefulness in Scenario
If we use ONNX, the problem we faced could be solved. Since ONNX makes it possible to have a compatible model between different frameworks, the most widely used two frameworks PyTorch and Tensorflow can be utilized. Also, if PyTorch model is converted as CoreML, the weight trained from the cloud can be used from devices such as iPhone or iPad.
Example
The problem we faced were the different frameworks of the models. So the example given are using the model trained from PyTorch on the Keras.
First, we trained a ResNet18 on PyTorch and saved it as in model.onnx format.
def train(dataloader, model, criterion, optimizer, num_epochs=20):
for epoch in range(num_epochs):
optimizer.step()
model.train()
running_loss = 0.0
running_corrects = 0
n = 0
for inputs, labels in dataloader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
.
.
.
args = torch.randn(1, 3, 32, 32, device=device)
ONNX_PATH = './ckpt/resnet18.onnx'
torch.onnx.export(model, args, ONNX_PATH, verbose=False)
Next, we are going to use the trained model on Keras using package onnx2keras.
import onnx
from onnx2keras import onnx_to_keras
# Load ONNX model
onnx_path = "ckpt/resnet18.onnx"
onnx_model = onnx.load(onnx_path)
keras_model = onnx_to_keras(onnx_model, ['input.1'])
keras_model.summary()
The following code results as below.
Model: "model"
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input.1 (InputLayer) [(None, 3, 32, 32)] 0 []
192_pad (ZeroPadding2D) (None, 3, 38, 38) 0 ['input.1[0][0]']
192 (Conv2D) (None, 64, 16, 16) 9472 ['192_pad[0][0]']
125 (Activation) (None, 64, 16, 16) 0 ['192[0][0]']
126_pad (ZeroPadding2D) (None, 64, 18, 18) 0 ['125[0][0]']
.
.
.
188 (Activation) (None, 512, 1, 1) 0 ['187[0][0]']
189 (GlobalAveragePooling2D) (None, 512) 0 ['188[0][0]']
189_EXPAND1 (Lambda) (None, 512, 1) 0 ['189[0][0]']
189_EXPAND2 (Lambda) (None, 512, 1, 1) 0 ['189_EXPAND1[0][0]']
190 (Reshape) (None, 512) 0 ['189_EXPAND2[0][0]']
191 (Dense) (None, 10) 5130 ['190[0][0]']
==================================================================================================
Total params: 11,176,842
Trainable params: 11,176,842
Non-trainable params: 0
__________________________________________________________________________________________________
At last, we converted the saved .onnx as Keras and run the code.
import onnx
from onnx2keras import onnx_to_keras
import keras
import numpy as np
# Load ONNX model
onnx_path = "ckpt/resnet18.onnx"
onnx_model = onnx.load(onnx_path)
keras_model = onnx_to_keras(onnx_model, ['input.1'])
.
.
.
y_pred = keras_model.predict(x_test)
y_pred = np.argmax(y_pred, axis=1)
acc = np.sum(y_pred == y_test.reshape(-1))/10000 * 100
print("Accuracy: %.2f%%" % (acc))
All the codes used in this post are available in https://github.com/sanglee325/onnx-pytorch
댓글남기기