Working with Other Frameworks

In Chapter 4, Working with Caffe, we learnt about Caffe and its relationship with Caffe2. We examined the Caffe and Caffe2 model file formats and looked at the process of importing a pre-trained Caffe model into Caffe2 using AlexNet as an example. In this chapter, we will look at how to export from, and import to, Caffe2 from other popular DL frameworks. And we will also look at how to enable other DL frameworks to use a model trained with Caffe2.

The topics covered in this chapter are as follows:

• The ONNX model format
• Support for ONNX in Caffe2
• How to export a Caffe2 model to ONNX format
• How to import an ONNX model into Caffe2
• How to visualize ONNX models

Open Neural Network Exchange

Open Neural Network Exchange (ONNX), typically pronounced as on-niks, is a format to represent a computation graph, with support for a wide variety of operators and data types. This format is general enough to support both neural networks and traditional ML models. Started by Facebook and Microsoft, this format has quickly gained a reputation as a popular format for the export and import of deep neural networks among most DL frameworks.

Installing ONNX

The ONNX source code can be found online at: https://github.com/onnx/onnx This includes definitions of the format and scripts to operate on ONNX files. Libraries and tools to convert from and to specific DL framework formats are usually provided...

ONNX in Caffe2

Caffe2 has built-in support for ONNX. This includes support for exporting Caffe2 models to ONNX format and importing ONNX models directly for inference in Caffe2. C++ source files related to Caffe2's support of ONNX can be found in the onnx directory in the Caffe2 source code. Python source files that provide the frontend and backend support for ONNX can be found in the python/onnx directory in the Caffe2 source code.

The onnx/onnx_exporter.h and onnx/onnx_exporter.cc contain the definitions necessary to export a Caffe2 model to ONNX format. Support for exporting from Caffe2 to ONNX includes details such as the mapping from Caffe2 to ONNX for operators, data types, and transformations of data.

For example, in onnx/onnx_exporter.cc we find the following mapping of some Caffe2 operators to ONNX operators:

const std::unordered_map<std::string, std::string>...

Exporting the Caffe2 model to ONNX

Caffe2 models can be easily exported to ONNX format using Python. This enables a vast number of other DL frameworks to use our Caffe2 models for training and inference. The frontend module provided by Caffe2-ONNX does all of the heavy lifting of the exporting. This module is located as the python/onnx/frontend.py file in the Caffe2 source code.

The ch5/export_to_onnx.py script provided along with this book's source code shows how to export an existing Caffe2 model to ONNX format. As an example, consider converting the Caffe2 model of AlexNet that we created in Chapter 4, Working with Caffe. We exported the operators and the weights of this network in Caffe2 to the files predict_net.pb and init_net.pb files respectively.

We can invoke the ONNX conversion script, as follows, to convert this Caffe2 model to an ONNX file named alexnet.onnx:

...

Using the ONNX model in Caffe2

In the previous section, we converted a Caffe2 model to ONNX format so that it could be used with other DL frameworks. In this section, we will learn how to use an ONNX model exported from other DL frameworks into Caffe2 for inference.

The backend module provided in the Caffe2 ONNX package enables this import of the ONNX model to Caffe2. This can be seen in the backend.py file in the python/onnx directory in the Caffe2 source code.

The ch5/run_onnx_model.py script provided along with this book's source code demonstrates how to load an ONNX model to Caffe2, and run an inference on an input image using that model.

The script first imports the Python modules necessary to work with the images (PIL.Image), Caffe2, and ONNX (caffe2.python.onnx.backend) as follows:

# Std
import PIL.Image
import json
import sys

# Ext
import numpy as np
from caffe2...

Visualizing the ONNX model

When working with ONNX models, it can be useful to have a tool that can help in visualizing the network structure. ONNX ships with such a script called net_drawer.py. You can find this tool in the onnx/onnx/tools directory in the ONNX source repository. If you installed ONNX from its Python package, then you can find this script at /usr/local/lib/python2.7/dist-packages/onnx/tools/net_drawer.py.

This script can be applied to convert an ONNX file to a directed acyclic graph representation of the network in the GraphViz DOT format. For example, consider the ONNX file alexnet.onnx that we obtained in the earlier section on converting from the Caffe2 model to the ONNX model.

We can convert this AlexNet ONNX file to a DOT file using the following command:

$ python /usr/local/lib/python2.7/dist-packages/onnx/tools/net_drawer.py --input alexnet.onnx --output...

Summary

In this chapter, we introduced the details of the ONNX format, a popular representation for DL models. We examined how it depicts the intermediate representation and operators. We then looked at support for ONNX in Caffe2. Using AlexNet as the example, we looked at how to convert a Caffe2 model file to ONNX format. We also looked at the reverse process: importing an ONNX model file into Caffe2, and then using it for inference. Finally, we looked at a useful tool to visualize the graph representation of an ONNX file.