|
|
|
|
|
|
|
|
"""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 |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
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() |
|
|
|
