|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +""" |
| 3 | +Caffe2 ONNX Primer |
| 4 | +================== |
| 5 | +**Author**: `Nathan Inkawhich <https://github.com/inkawhich>`_ |
| 6 | +
|
| 7 | +This tutorial is a brief look at how to use Caffe2 and |
| 8 | +`ONNX <http://onnx.ai/about>`_ together. More specifically, we will |
| 9 | +show how to export a model from Caffe2 to ONNX and how to import a model |
| 10 | +from ONNX into Caffe2. Hopefully, the motivation is clear but this |
| 11 | +tutorial shows how to use the very fast and efficient Caffe2 framework |
| 12 | +with the flexibility enabling ONNX framework. One important fact to keep |
| 13 | +in mind is that ONNX is designed to enable deployment and *inference* in |
| 14 | +frameworks other than where the model was trained. Currently, there is |
| 15 | +no streamlined way to finetune ONNX models. The workflow for this |
| 16 | +document is as follows: |
| 17 | +
|
| 18 | +- Run prediction with a Caffe2 model and collect initial prediction |
| 19 | +- Export the Caffe2 model to ONNX format |
| 20 | +- Import the saved ONNX model back into Caffe2 |
| 21 | +- Run prediction on imported model and verify results |
| 22 | +
|
| 23 | +Let's get started with some imports. |
| 24 | +
|
| 25 | +""" |
| 26 | + |
| 27 | +from __future__ import absolute_import |
| 28 | +from __future__ import division |
| 29 | +from __future__ import print_function |
| 30 | +from __future__ import unicode_literals |
| 31 | +import numpy as np |
| 32 | +import operator |
| 33 | +from caffe2.proto import caffe2_pb2 |
| 34 | +from caffe2.python import core, workspace, models |
| 35 | +import onnx |
| 36 | +import caffe2.python.onnx.frontend # Required for Caffe2->ONNX export |
| 37 | +import caffe2.python.onnx.backend # Required for ONNX->Caffe2 import |
| 38 | + |
| 39 | + |
| 40 | +###################################################################### |
| 41 | +# Inputs |
| 42 | +# ------ |
| 43 | +# |
| 44 | +# Now we will specify the inputs. The *MODELS\_DIR* is where the |
| 45 | +# downloaded Caffe2 models are saved, the *MODEL\_NAME* is the name of the |
| 46 | +# model we want to use, and *SIZE* is the size of image the model expects. |
| 47 | +# For more information about downloading a pretrained Caffe2 model, see |
| 48 | +# the `Loading Pretrained Models |
| 49 | +# Tutorial <https://github.com/caffe2/tutorials/blob/master/Loading_Pretrained_Models.ipynb>`__. |
| 50 | +# |
| 51 | + |
| 52 | +# User Inputs |
| 53 | +MODELS_DIR = "../models" |
| 54 | +MODEL_NAME = "squeezenet" # e.g. [squeezenet, bvlc_alexnet, bvlc_googlenet, bvlc_reference_caffenet] |
| 55 | +SIZE = 224 |
| 56 | + |
| 57 | +# Construct path strings from inputs |
| 58 | +INIT_NET = "{}/{}/init_net.pb".format(MODELS_DIR, MODEL_NAME) |
| 59 | +PREDICT_NET = "{}/{}/predict_net.pb".format(MODELS_DIR, MODEL_NAME) |
| 60 | +ONNX_MODEL = "{}/{}/my_model.onnx".format(MODELS_DIR, MODEL_NAME) # we will create this |
| 61 | + |
| 62 | + |
| 63 | +###################################################################### |
| 64 | +# Load Caffe2 Model |
| 65 | +# ----------------- |
| 66 | +# |
| 67 | +# Before we perform the export we will first load the pretrained init and |
| 68 | +# predict nets, then create a *Predictor*. Next, we will create a random |
| 69 | +# input to get a baseline result for comparision later. Take note of the |
| 70 | +# predicted label and confidence. |
| 71 | +# |
| 72 | + |
| 73 | +# Generate random NCHW input to run model |
| 74 | +# This is a placeholder for any real image that is processed and |
| 75 | +# put in NCHW order. |
| 76 | +image = np.random.rand(1,3,SIZE,SIZE).astype(np.float32) |
| 77 | +print("Input Shape: ",image.shape) |
| 78 | + |
| 79 | +# Prepare the nets |
| 80 | +predict_net = caffe2_pb2.NetDef() |
| 81 | +with open(PREDICT_NET, 'rb') as f: |
| 82 | + predict_net.ParseFromString(f.read()) |
| 83 | +init_net = caffe2_pb2.NetDef() |
| 84 | +with open(INIT_NET, 'rb') as f: |
| 85 | + init_net.ParseFromString(f.read()) |
| 86 | + |
| 87 | +# Initialize the predictor from the nets |
| 88 | +p = workspace.Predictor(init_net, predict_net) |
| 89 | + |
| 90 | +#### Run the sample data |
| 91 | + |
| 92 | +# Run the net and return prediction |
| 93 | +results = p.run({'data': image}) |
| 94 | +results = np.asarray(results) |
| 95 | +print("Results Shape: ", results.shape) |
| 96 | + |
| 97 | +# Quick way to get the top-1 prediction result |
| 98 | +curr_pred, curr_conf = max(enumerate(np.squeeze(results)), key=operator.itemgetter(1)) |
| 99 | +print("Top-1 Prediction: {} @ {}".format(curr_pred, curr_conf)) |
| 100 | + |
| 101 | + |
| 102 | + |
| 103 | +###################################################################### |
| 104 | +# Caffe2 :math:`\rightarrow` ONNX Export |
| 105 | +# -------------------------------------- |
| 106 | +# |
| 107 | +# Finally, we have reached the interesting stuff. It is not hard to |
| 108 | +# imagine why one may want to export a Caffe2 model to ONNX. Maybe you |
| 109 | +# have a cool idea for an iPhone app and want to use a model trained in |
| 110 | +# Caffe2 with CoreML as part of the app. Or, maybe you have a system built |
| 111 | +# in Tensorflow but want to test out a model from the Caffe2 Model Zoo. |
| 112 | +# ONNX enables this interoperability by allowing models to be imported and |
| 113 | +# exported into different frameworks (for inference!). |
| 114 | +# |
| 115 | +# The code below shows how to **export** a model trained in Caffe2 to ONNX |
| 116 | +# format. Once in ONNX format, the model can be imported into any other |
| 117 | +# compatible framework to be used for *inference*. From the Caffe2 side, |
| 118 | +# we only need the previously loaded *init\_net* and *predict\_net* |
| 119 | +# *caffe2\_pb2.NetDef* objects. |
| 120 | +# |
| 121 | +# There are only a few steps to export once the nets are loaded. First, we |
| 122 | +# must declare (via Python dictionary) the type and shape of inputs and |
| 123 | +# outputs of the model. This information is not explicitly specified in |
| 124 | +# the Caffe2 model architecture but is required by ONNX. Next, we must |
| 125 | +# make sure the model has a name, otherwise the internal model checks in |
| 126 | +# the ONNX converter will fail. Then, all thats left to do is create the |
| 127 | +# ONNX model, check it, and save it. |
| 128 | +# |
| 129 | + |
| 130 | +# We need to provide type and shape of the model inputs |
| 131 | +data_type = onnx.TensorProto.FLOAT |
| 132 | +data_shape = (1, 3, 224, 224) |
| 133 | +value_info = { |
| 134 | + 'data': (data_type, data_shape) |
| 135 | +} |
| 136 | + |
| 137 | +# Make sure the net has a name. Otherwise, the checker will fail. |
| 138 | +if predict_net.name == "": |
| 139 | + predict_net.name = "ModelNameHere" |
| 140 | + |
| 141 | +# Create the ONNX model |
| 142 | +onnx_model = caffe2.python.onnx.frontend.caffe2_net_to_onnx_model( |
| 143 | + predict_net, |
| 144 | + init_net, |
| 145 | + value_info, |
| 146 | +) |
| 147 | + |
| 148 | +# Check the ONNX model. Exception will be thrown if there is a problem here. |
| 149 | +onnx.checker.check_model(onnx_model) |
| 150 | + |
| 151 | +# Save the ONNX model |
| 152 | +onnx.save(onnx_model, ONNX_MODEL) |
| 153 | + |
| 154 | + |
| 155 | +###################################################################### |
| 156 | +# ONNX :math:`\rightarrow` Caffe2 Import |
| 157 | +# -------------------------------------- |
| 158 | +# |
| 159 | +# Now suppose someone has trained Alexnet2.0 which gets 99.9% top-1 test |
| 160 | +# accuracy on ImageNet ... *gasp* ... in Tensorflow. As a Caffe2 user, all |
| 161 | +# we have to do is convince them to convert the model to ONNX format, then |
| 162 | +# we can import it and use it. Since we are running out of time in this |
| 163 | +# 5-minute primer, here we will only show how to import the model we just |
| 164 | +# exported back into Caffe2. The import happens in a single load command |
| 165 | +# (``onnx.load``), then we can start feeding the model data in just one |
| 166 | +# more command (``run_model``). Also, note that the predictions from this |
| 167 | +# imported model and the original model are the exact same, indicating |
| 168 | +# nothing was lost in the export/import process. |
| 169 | +# |
| 170 | + |
| 171 | +# Load the ONNX model |
| 172 | +model = onnx.load(ONNX_MODEL) |
| 173 | + |
| 174 | +# Run the ONNX model with Caffe2 |
| 175 | +outputs = caffe2.python.onnx.backend.run_model(model, [image]) |
| 176 | +print("Output Shape: ", np.array(outputs).shape) |
| 177 | + |
| 178 | +# Get model prediction |
| 179 | +curr_pred, curr_conf = max(enumerate(np.squeeze(results)), key=operator.itemgetter(1)) |
| 180 | +print("Top-1 Prediction: {} @ {}".format(curr_pred, curr_conf)) |
| 181 | + |
| 182 | + |
| 183 | + |
| 184 | +###################################################################### |
| 185 | +# Hopefully it is clear that the caffe2-onnx interface for both importing |
| 186 | +# and exporting is relatively simple. For more information about ONNX and |
| 187 | +# to see more tutorials on using ONNX with different frameworks see the |
| 188 | +# `ONNX Tutorials <https://github.com/onnx/tutorials>`__. Also, although |
| 189 | +# importing and exporting with Caffe2 is supported, and exporting a model |
| 190 | +# from PyTorch to ONNX is supported, *importing* an ONNX model into |
| 191 | +# PyTorch is *NOT*, but is coming soon! |
| 192 | +# |
| 193 | +# Here are some more cool ONNX resources for the curious reader: |
| 194 | +# |
| 195 | +# - `ONNX Python API |
| 196 | +# Overview <https://github.com/onnx/onnx/blob/master/docs/PythonAPIOverview.md>`__ |
| 197 | +# - `ONNX Model Zoo <https://github.com/onnx/models>`__ |
| 198 | +# - `ONNX |
| 199 | +# Operators <https://github.com/onnx/onnx/blob/master/docs/Operators.md>`__ |
| 200 | +# - `ONNX Tutorials <https://github.com/onnx/tutorials>`__ |
| 201 | +# |
0 commit comments