Caffe - 网络结构 prototxt 可视化

Netscope 是对于 Caffe 网络结构可视化的很好的工具.

Caffe 也有基于 Python 的网络结构可视化脚本,如:

#!/usr/bin/env python
#! --*-- coding:utf-8 --*--
"""
Caffe prototxt 网络结构可视化
"""
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from google.protobuf import text_format

import sys
sys.path.insert(0, '/path/to/caffe/python')
import caffe
import caffe.draw
from caffe.proto import caffe_pb2


def parse_args():
    """
    输入参数
    """

    parser = ArgumentParser(description=__doc__,
                            formatter_class=ArgumentDefaultsHelpFormatter)
    
    # Caffe prototxt 网络定义文件
    parser.add_argument('--prototxt', 
                        help='Input network prototxt file')
    # prototxt 训练或测试结构图
    parser.add_argument('--phase',
                        help='TRAIN, or TEST, or ALL',
                        default="ALL")
    # 输出保存的网络结构图
    parser.add_argument('--output_image', 
                        default='output_image.png',
                        help='Output image file')
    # 网络结构的可视化方式,可选 'TB' 'LR'等.
    parser.add_argument('--vis_order',
                        help='TB(top-bottom) or RL(right-left)',
                        default='LR')
    # 可视化待学习网络层的学习率lr,multplier,
    # 如 Convolution, Deconvolution, InnerProduct 等网络层
    parser.add_argument('--display_lrm', action='store_true',
                       default=False)

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    print('[INFO] Prototxt %s' % args.prototxt)
    print('[INFO] Net phase %s' % args.phase)
    print('[INFO] Vis order %s' % args.vis_order)
    print('[INFO] Display lr and multplier %s' % args.display_lrm)
    print('[INFO] Output image %s' % args.output_image)
    
    assert args.phase in ['TRAIN', 'TEST', 'ALL']
    
    net = caffe_pb2.NetParameter()
    text_format.Merge(open(args.prototxt).read(), net)
    print('[INFO] Drawing net to %s' % args.prototxt)
    
    caffe.draw.draw_net_to_file(net, 
                                args.output_image, 
                                args.vis_order,
                                args.phase, 
                                args.display_lrm)
    print('[INFO] Done.')
   

if __name__ == '__main__':
    main()

输出结果如:

Last modification:December 1st, 2018 at 04:01 pm

Leave a Comment