Installing Caffe2 with CUDA in Conda

Deprecation warning

Since May 2008, Caffe2 has been merged in PyTorch. To install the lastest version of Caffe2, simply get PyTorch. The instructions for installing PyTorch can be accessed here.

The following guide is kept here for posterity.

---

The following guide shows you how to install install caffe2 with CUDA under Conda virtual environment. This guide is meant for machines running on Ubuntu 16.04 equipped with NVIDIA GPUs with CUDA support. i.e it assumes CUDA is already installed by a system admin.

Assumptions

• Ubuntu OS
• NVIDIA GPU
• Conda (see installation instructions here)
• CUDA (installed by system admin)

Specifications

This guide is written for the following specs:

• Ubuntu 16.04
• Python 2.7
• CUDA 8
• cuDNN v7.1
• Miniconda 2
• OpenCV3

Guide

First, get cuDNN by following this cuDNN Guide.

Let’s create a virtual Conda environment called “caffe2”:

conda create -n caffe2 python=2.7

You many of course use a different environment name, just be sure to adjust accordingly for the rest of this guide

After it prepares the environment and installs the default packages, activate the virtual environment via:

conda activate caffe2
# to deactivate: conda deactivate caffe2

Now let’s install the necessary dependencies in our current caffe2 environment:

conda install future gflags glog lmdb mkl mkl-include numpy opencv protobuf snappy six cmake -y

Let’s clone caffe2’s repo and its submodules into our home directory.

cd ~
git clone --recursive https://github.com/caffe2/caffe2.git
cd caffe2
git submodule update --init --recursive

We shall avoid polluting the caffe2 source tree by building within a build folder

mkdir build && cd build
pwd #=> ~/caffe2/build

We shall now build the package using CMake with the following flags

cmake -DCUDNN_INCLUDE_DIR=$CUDA_HOME/include -DCUDNN_LIBRARY=$CUDA_HOME/lib64/libcudnn.so -DCMAKE_PREFIX_PATH=$CONDA_PREFIX -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX ..

• CMake variable CUDNN_INCLUDE_DIR indicates where to find the include directory for your cuDNN
• CMake variable CUDNN_LIBRARY indicates where to find the libcudnn.so for your cuDNN
• CMake variable CMAKE_PREFIX_PATH tells CMake to look for packages in your conda environment before looking in system install locations (like /usr/local)
• CMake variable CMAKE_INSTALL_PREFIX indicates where to install Caffe2 binaries such as libcaffe2.dylib after Caffe2 has been successfully built, the default is /usr/local which will require administrator privilege
• CMake variable CPU indicates whether or not to use CPU=only installation

Let’s find out how many cores your machine has

cat /proc/cpuinfo | grep processor | wc -l
#=> n

Let’s make the package efficiently by maximising the number of jobs for it. General rule of thumb is to use 1 + n number of jobs where n is the output from the previous command. i.e. number of cores. Mine was 24 so I run the following

make -j 25

After make is completed, we are now finally ready to install

make install

You’d think we’re done, but not quite! We have to point the $PYTHONPATH environment variable to our build folder like so

export PYTHONPATH=$HOME/caffe2/build:$PYTHONPATH

However it will be tedious to type that everytime we activate our environment. You may append that line to .bash_profile or .bashrc but some variables such as $PYTHONPATH are potentially used in many environments and it could lead to python import errors when the paths contain different modules sharing the same name. For instance, both caffe and caffe2 contain a module named ‘caffe’.

The solution to overcome this is to write a script to save our environment variables within our environemnt so that they get loaded automatically every time we activate our environment and get unset automatically when we deactivate our environment. The following steps are an adaptation of this guide stated in the official Conda documentation.

Let’s enter our environment directory and do the following

cd $CONDA_PREFIX
mkdir -p ./etc/conda/activate.d
mkdir -p ./etc/conda/deactivate.d
touch ./etc/conda/activate.d/env_vars.sh
touch ./etc/conda/deactivate.d/env_vars.sh

Edit ./etc/conda/activate.d/env_vars.sh as follows:

#!/bin/sh

export PYTHONPATH=$HOME/caffe2/build:$PYTHONPATH

Edit ./etc/conda/deactivate.d/env_vars.sh as follows:

#!/bin/sh

unset PYTHONPATH

Now let’s reload the current environment to reflect the variables

conda activate caffe2

We are now ready to test if caffe2 has installed correctly

cd ~
# To check if Caffe2 build was successful
python -c 'from caffe2.python import core' 2>/dev/null && echo "Success" || echo "Failure"
#=> Success

# To check if Caffe2 GPU build was successful
# This must print a number > 0 in order to use Detectron
python2 -c 'from caffe2.python import workspace; print(workspace.NumCudaDevices())'
#=> 2