Creating a diagram for a caffe neural network architecture

In this tutorial, I will explain how to save an image with the diagram of a caffe neural network architecture The Python notebook with the code is in this url.

If we want to create a visualisation for the caffe neural network, we can use the Python library. The first requirement is to install GraphViz an pydot. You can find instructions here for GraphViz.

Now, these are the Python imports required:

import sys

from google.protobuf import text_format

caffe_library = 'caffelibrary/'
sys.path.insert(0, caffe_library)
import caffe;
import caffe.draw

This is the method we can use to create the diagram:

def create_network_architecture_diagram(networkfilepath, fileextextension='png'):
    with open(networkfilepath, "r") as f:
        net_architecture_text = f.read()

    net_architecture = caffe.proto.caffe_pb2.NetParameter()
    text_format.Merge(net_architecture_text, net_architecture)
    caffe.draw.draw_net_to_file(net_architecture, networkfilepath + '.' + fileextextension, "LR")

We can optionally select the format of the output, any file extension supported by GraphViz is valid, e.g.: png, jpg, svg, eps. An example of how to use it is shown here:

create_network_architecture_diagram("Path to the prototxy network file", "eps")

The following diagram corresponds to the architecture of the network "net surgery" provided in the caffe examples: