Windows下利用Python3绘制Caffe中的网络结构

在Windows10+Anaconda3环境下利用Caffe中的draw_net.py文件绘制Caffe中的网络结构图

1.必要安装

第一步：安装pydotplus。在cmd命令窗口输入：pip install pydotplus。
第二步：安装Graphviz-2.38。在http://www.graphviz.org/Download_windows.php页面中下载Graphviz安装包，并完成安装（例如安装目录：C:/Program Files (x86)/Graphviz2.38/bin/）；右击“此电脑”打开“属性”，进入系统变量，添加C:/Program Files (x86)/Graphviz2.38/bin/到path中。

2.修改draw_net.py和draw.py文件

（1）打开caffe的根目录下python/caffe/draw.py文件，修改get_layer_label函数中的elif layer.type == 'Pooling':以下的内容，修改如下：

def get_layer_label(layer, rankdir):
    """Define node label based on layer type.

    Parameters
    ----------
    layer : ?
    rankdir : {'LR', 'RL', 'TB', 'BT'}
        Direction of graph layout.

    Returns
    -------
    string :
        A label for the current layer
    """

    if rankdir in ('TB', 'BT'):
        # If graph orientation is vertical, horizontal space is free and
        # vertical space is not; separate words with spaces
        separator = ' '
    else:
        # If graph orientation is horizontal, vertical space is free and
        # horizontal space is not; separate words with newlines
        separator = '\\n'

    if layer.type == 'Convolution' or layer.type == 'Deconvolution':
        # Outer double quotes needed or else colon characters don't parse
        # properly
        node_label = '"%s%s(%s)%skernel size: %d%sstride: %d%spad: %d"' %\
                     (layer.name,
                      separator,
                      layer.type,
                      separator,
                      layer.convolution_param.kernel_size[0] if len(layer.convolution_param.kernel_size._values) else 1,
                      separator,
                      layer.convolution_param.stride[0] if len(layer.convolution_param.stride._values) else 1,
                      separator,
                      layer.convolution_param.pad[0] if len(layer.convolution_param.pad._values) else 0)
    elif layer.type == 'Pooling':
        pooling_types_dict = get_pooling_types_dict()
        node_label = '"%s%s(%s %s)%skernel size: %d%sstride: %d%spad: %d"' %\
                     (layer.name,
                      separator,
                      pooling_types_dict[layer.pooling_param.pool],
                      layer.type,
                      separator,
                      layer.pooling_param.kernel_size,
                      separator,
                      layer.pooling_param.stride,
                      separator,
                      layer.pooling_param.pad)
    else:
        node_label = '"%s%s(%s)"' % (layer.name, separator, layer.type)
    return node_label

（2）若在pycharm下不使用命令窗口方式，则编辑caffe的根目录下python/draw_net.py文件，修改main函数如下：

def main():
    """
    input_net_proto_file: input network prototxt file.
    output_image_file: output image file
    phase: Which network phase to draw: can be TRAIN, TEST, or ALL. Default=ALL.
           If ALL, then all layers are drawn regardless of phase.
    rankdir: One of TB (top-bottom, i.e., vertical)、BT(bottom-top, i.e., vertical), 
             LR (left-right, i.e., horizontal), RL (right-left, i.e., horizontal);
    """
    # the file of the net prototxt
    input_net_proto_path = '../examples/mnist/'
    input_net_proto_name = 'lenet_train_test.prototxt'
    input_net_proto_file = input_net_proto_path + input_net_proto_name

    net = caffe_pb2.NetParameter()
    text_format.Merge(open(input_net_proto_file).read(), net)
    rankdir = 'LR'
    # draw the TRAIN or TEST phase network
    phase = "ALL"
    if phase == "TRAIN":
        phase = caffe.TRAIN
    elif phase == "TEST":
        phase = caffe.TEST
    elif phase == "ALL":
        phase = None

    # the path to a file where the networks visualization will be stored
    output_net_proto_path = '../examples/mnist/'
    output_net_proto_file = output_net_proto_path + 'Lenet_' + str(phase) + '.jpg'
    draw.draw_net_to_file(net, output_net_proto_file, rankdir, phase)

**input_net_proto_file：指prototxt文件
output_image_file：指保存后的图片名称和格式**

若在pycharm下使用命令窗口方式，则在下面界面中

输入命令：python ./python/draw_net.py ./examples/mnist/lenet_train_test.prototxt lenet.jpg
也可以指定图片的保存位置：python ./python/draw_net.py ./examples/mnist/lenet_train_test.prototxt ./examples/mnist/lenet.jpg，图片保存在examples/mnist/目录下。
在上述命令中–rankdir默认值为LR，–phase默认值为ALL。