How to use caffe inside python virtual environment ?

In the previous blog post, we installed caffe for python3 with GPU support. But when we try to import caffe inside a python virtual environment, we may get some error saying caffe is not available. Take a look at the error below.

Python 3.5.2 (default, Nov 12 2018, 13:43:14) 
[GCC 5.4.0 20160609] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import caffe
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: No module named 'caffe'
>>> 

We can successfully import caffe outside the virtual environment, but it is not the same inside the virtual environment.

Solution

We need to set the environment variables once we are inside the virtual environment. There are a couple of ways to do it.

I. Set the environment variables manually inside the virtual environment. In the terminal after your virtual environment is activated, key in the following commands.

export LD_LIBRARY_PATH=/path/to/caffe/install/lib:$LD_LIBRARY_PATH
export PYTHONPATH=/path/to/caffe/install/python:$PYTHONPATH
export CAFFE_ROOT=/path/to/caffe/install/

With this approach, we must run the above commands everytime we get into the virtual environment.

II. Another way is to set the environment variables in activate script of the virtual environment. For this we need to add the above commands to the activate file of your virtual environment.

$ cd /your/virtual/environment/path
$ cd bin
$ echo "# Caffe
export LD_LIBRARY_PATH=/path/to/caffe/install/lib:$LD_LIBRARY_PATH
export PYTHONPATH=/path/to/caffe/install/python:$PYTHONPATH
export CAFFE_ROOT=/path/to/caffe/install/" >> activate

Now everytime you activate your virtual environment you’ll be able to import caffe without any issues.