Converting models from PyTorch to Caffe2 using ONNX

Sunday, February 18, 2018

This tutorial describes how to use ONNX to convert a model defined in PyTorch into the ONNX format and then convert it into Caffe2. However, if you follow the way in the tutorial to install onnx, onnx-caffe2 and Caffe2, you may experience some errors. Here I provide a solution to solve this problem.

Install CUDA-9.0

Follow Install Nvidia CUDA-9.0 on Ubuntu 16.04 for Deep Learning to install CUDA-9.0.

Install Anaconda

Download Anaconda from here. Choose Python 3.6 version 64-BIT INSTALLER to install it, then update it:

conda update conda
conda create -n pytorch
source activate pytorch

Install Pytorch and Caffe2

conda install pytorch torchvision cuda90 -c pytorch
conda install -c caffe2 caffe2-cuda9.0-cudnn7

Install ONNX and ONNX-Caffe2

pip install onnx onnx-caffe2
conda install matplotlib scikit-image pydot