141 lines
4.8 KiB
Python
141 lines
4.8 KiB
Python
#!/usr/bin/env python
|
|
|
|
"""Net summarization tool.
|
|
|
|
This tool summarizes the structure of a net in a concise but comprehensive
|
|
tabular listing, taking a prototxt file as input.
|
|
|
|
Use this tool to check at a glance that the computation you've specified is the
|
|
computation you expect.
|
|
"""
|
|
|
|
from caffe.proto import caffe_pb2
|
|
from google import protobuf
|
|
import re
|
|
import argparse
|
|
|
|
# ANSI codes for coloring blobs (used cyclically)
|
|
COLORS = ['92', '93', '94', '95', '97', '96', '42', '43;30', '100',
|
|
'444', '103;30', '107;30']
|
|
DISCONNECTED_COLOR = '41'
|
|
|
|
def read_net(filename):
|
|
net = caffe_pb2.NetParameter()
|
|
with open(filename) as f:
|
|
protobuf.text_format.Parse(f.read(), net)
|
|
return net
|
|
|
|
def format_param(param):
|
|
out = []
|
|
if len(param.name) > 0:
|
|
out.append(param.name)
|
|
if param.lr_mult != 1:
|
|
out.append('x{}'.format(param.lr_mult))
|
|
if param.decay_mult != 1:
|
|
out.append('Dx{}'.format(param.decay_mult))
|
|
return ' '.join(out)
|
|
|
|
def printed_len(s):
|
|
return len(re.sub(r'\033\[[\d;]+m', '', s))
|
|
|
|
def print_table(table, max_width):
|
|
"""Print a simple nicely-aligned table.
|
|
|
|
table must be a list of (equal-length) lists. Columns are space-separated,
|
|
and as narrow as possible, but no wider than max_width. Text may overflow
|
|
columns; note that unlike string.format, this will not affect subsequent
|
|
columns, if possible."""
|
|
|
|
max_widths = [max_width] * len(table[0])
|
|
column_widths = [max(printed_len(row[j]) + 1 for row in table)
|
|
for j in range(len(table[0]))]
|
|
column_widths = [min(w, max_w) for w, max_w in zip(column_widths, max_widths)]
|
|
|
|
for row in table:
|
|
row_str = ''
|
|
right_col = 0
|
|
for cell, width in zip(row, column_widths):
|
|
right_col += width
|
|
row_str += cell + ' '
|
|
row_str += ' ' * max(right_col - printed_len(row_str), 0)
|
|
print row_str
|
|
|
|
def summarize_net(net):
|
|
disconnected_tops = set()
|
|
for lr in net.layer:
|
|
disconnected_tops |= set(lr.top)
|
|
disconnected_tops -= set(lr.bottom)
|
|
|
|
table = []
|
|
colors = {}
|
|
for lr in net.layer:
|
|
tops = []
|
|
for ind, top in enumerate(lr.top):
|
|
color = colors.setdefault(top, COLORS[len(colors) % len(COLORS)])
|
|
if top in disconnected_tops:
|
|
top = '\033[1;4m' + top
|
|
if len(lr.loss_weight) > 0:
|
|
top = '{} * {}'.format(lr.loss_weight[ind], top)
|
|
tops.append('\033[{}m{}\033[0m'.format(color, top))
|
|
top_str = ', '.join(tops)
|
|
|
|
bottoms = []
|
|
for bottom in lr.bottom:
|
|
color = colors.get(bottom, DISCONNECTED_COLOR)
|
|
bottoms.append('\033[{}m{}\033[0m'.format(color, bottom))
|
|
bottom_str = ', '.join(bottoms)
|
|
|
|
if lr.type == 'Python':
|
|
type_str = lr.python_param.module + '.' + lr.python_param.layer
|
|
else:
|
|
type_str = lr.type
|
|
|
|
# Summarize conv/pool parameters.
|
|
# TODO support rectangular/ND parameters
|
|
conv_param = lr.convolution_param
|
|
if (lr.type in ['Convolution', 'Deconvolution']
|
|
and len(conv_param.kernel_size) == 1):
|
|
arg_str = str(conv_param.kernel_size[0])
|
|
if len(conv_param.stride) > 0 and conv_param.stride[0] != 1:
|
|
arg_str += '/' + str(conv_param.stride[0])
|
|
if len(conv_param.pad) > 0 and conv_param.pad[0] != 0:
|
|
arg_str += '+' + str(conv_param.pad[0])
|
|
arg_str += ' ' + str(conv_param.num_output)
|
|
if conv_param.group != 1:
|
|
arg_str += '/' + str(conv_param.group)
|
|
elif lr.type == 'Pooling':
|
|
arg_str = str(lr.pooling_param.kernel_size)
|
|
if lr.pooling_param.stride != 1:
|
|
arg_str += '/' + str(lr.pooling_param.stride)
|
|
if lr.pooling_param.pad != 0:
|
|
arg_str += '+' + str(lr.pooling_param.pad)
|
|
else:
|
|
arg_str = ''
|
|
|
|
if len(lr.param) > 0:
|
|
param_strs = map(format_param, lr.param)
|
|
if max(map(len, param_strs)) > 0:
|
|
param_str = '({})'.format(', '.join(param_strs))
|
|
else:
|
|
param_str = ''
|
|
else:
|
|
param_str = ''
|
|
|
|
table.append([lr.name, type_str, param_str, bottom_str, '->', top_str,
|
|
arg_str])
|
|
return table
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Print a concise summary of net computation.")
|
|
parser.add_argument('filename', help='net prototxt file to summarize')
|
|
parser.add_argument('-w', '--max-width', help='maximum field width',
|
|
type=int, default=30)
|
|
args = parser.parse_args()
|
|
|
|
net = read_net(args.filename)
|
|
table = summarize_net(net)
|
|
print_table(table, max_width=args.max_width)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|