258 lines
10 KiB
Python
258 lines
10 KiB
Python
import torch
|
|
from .torch_utils import convert2cpu
|
|
|
|
def parse_cfg(cfgfile):
|
|
blocks = []
|
|
fp = open(cfgfile, 'r')
|
|
block = None
|
|
line = fp.readline()
|
|
while line != '':
|
|
line = line.rstrip()
|
|
if line == '' or line[0] == '#':
|
|
line = fp.readline()
|
|
continue
|
|
elif line[0] == '[':
|
|
if block:
|
|
blocks.append(block)
|
|
block = dict()
|
|
block['type'] = line.lstrip('[').rstrip(']')
|
|
# set default value
|
|
if block['type'] == 'convolutional':
|
|
block['batch_normalize'] = 0
|
|
else:
|
|
key, value = line.split('=')
|
|
key = key.strip()
|
|
if key == 'type':
|
|
key = '_type'
|
|
value = value.strip()
|
|
block[key] = value
|
|
line = fp.readline()
|
|
|
|
if block:
|
|
blocks.append(block)
|
|
fp.close()
|
|
return blocks
|
|
|
|
|
|
def print_cfg(blocks):
|
|
print('layer filters size input output');
|
|
prev_width = 416
|
|
prev_height = 416
|
|
prev_filters = 3
|
|
out_filters = []
|
|
out_widths = []
|
|
out_heights = []
|
|
ind = -2
|
|
for block in blocks:
|
|
ind = ind + 1
|
|
if block['type'] == 'net':
|
|
prev_width = int(block['width'])
|
|
prev_height = int(block['height'])
|
|
continue
|
|
elif block['type'] == 'convolutional':
|
|
filters = int(block['filters'])
|
|
kernel_size = int(block['size'])
|
|
stride = int(block['stride'])
|
|
is_pad = int(block['pad'])
|
|
pad = (kernel_size - 1) // 2 if is_pad else 0
|
|
width = (prev_width + 2 * pad - kernel_size) // stride + 1
|
|
height = (prev_height + 2 * pad - kernel_size) // stride + 1
|
|
print('%5d %-6s %4d %d x %d / %d %3d x %3d x%4d -> %3d x %3d x%4d' % (
|
|
ind, 'conv', filters, kernel_size, kernel_size, stride, prev_width, prev_height, prev_filters, width,
|
|
height, filters))
|
|
prev_width = width
|
|
prev_height = height
|
|
prev_filters = filters
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'maxpool':
|
|
pool_size = int(block['size'])
|
|
stride = int(block['stride'])
|
|
width = prev_width // stride
|
|
height = prev_height // stride
|
|
print('%5d %-6s %d x %d / %d %3d x %3d x%4d -> %3d x %3d x%4d' % (
|
|
ind, 'max', pool_size, pool_size, stride, prev_width, prev_height, prev_filters, width, height,
|
|
filters))
|
|
prev_width = width
|
|
prev_height = height
|
|
prev_filters = filters
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'avgpool':
|
|
width = 1
|
|
height = 1
|
|
print('%5d %-6s %3d x %3d x%4d -> %3d' % (
|
|
ind, 'avg', prev_width, prev_height, prev_filters, prev_filters))
|
|
prev_width = width
|
|
prev_height = height
|
|
prev_filters = filters
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'softmax':
|
|
print('%5d %-6s -> %3d' % (ind, 'softmax', prev_filters))
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'cost':
|
|
print('%5d %-6s -> %3d' % (ind, 'cost', prev_filters))
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'reorg':
|
|
stride = int(block['stride'])
|
|
filters = stride * stride * prev_filters
|
|
width = prev_width // stride
|
|
height = prev_height // stride
|
|
print('%5d %-6s / %d %3d x %3d x%4d -> %3d x %3d x%4d' % (
|
|
ind, 'reorg', stride, prev_width, prev_height, prev_filters, width, height, filters))
|
|
prev_width = width
|
|
prev_height = height
|
|
prev_filters = filters
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'upsample':
|
|
stride = int(block['stride'])
|
|
filters = prev_filters
|
|
width = prev_width * stride
|
|
height = prev_height * stride
|
|
print('%5d %-6s * %d %3d x %3d x%4d -> %3d x %3d x%4d' % (
|
|
ind, 'upsample', stride, prev_width, prev_height, prev_filters, width, height, filters))
|
|
prev_width = width
|
|
prev_height = height
|
|
prev_filters = filters
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'route':
|
|
layers = block['layers'].split(',')
|
|
layers = [int(i) if int(i) > 0 else int(i) + ind for i in layers]
|
|
if len(layers) == 1:
|
|
print('%5d %-6s %d' % (ind, 'route', layers[0]))
|
|
prev_width = out_widths[layers[0]]
|
|
prev_height = out_heights[layers[0]]
|
|
prev_filters = out_filters[layers[0]]
|
|
elif len(layers) == 2:
|
|
print('%5d %-6s %d %d' % (ind, 'route', layers[0], layers[1]))
|
|
prev_width = out_widths[layers[0]]
|
|
prev_height = out_heights[layers[0]]
|
|
assert (prev_width == out_widths[layers[1]])
|
|
assert (prev_height == out_heights[layers[1]])
|
|
prev_filters = out_filters[layers[0]] + out_filters[layers[1]]
|
|
elif len(layers) == 4:
|
|
print('%5d %-6s %d %d %d %d' % (ind, 'route', layers[0], layers[1], layers[2], layers[3]))
|
|
prev_width = out_widths[layers[0]]
|
|
prev_height = out_heights[layers[0]]
|
|
assert (prev_width == out_widths[layers[1]] == out_widths[layers[2]] == out_widths[layers[3]])
|
|
assert (prev_height == out_heights[layers[1]] == out_heights[layers[2]] == out_heights[layers[3]])
|
|
prev_filters = out_filters[layers[0]] + out_filters[layers[1]] + out_filters[layers[2]] + out_filters[
|
|
layers[3]]
|
|
else:
|
|
print("route error !!! {} {} {}".format(sys._getframe().f_code.co_filename,
|
|
sys._getframe().f_code.co_name, sys._getframe().f_lineno))
|
|
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] in ['region', 'yolo']:
|
|
print('%5d %-6s' % (ind, 'detection'))
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'shortcut':
|
|
from_id = int(block['from'])
|
|
from_id = from_id if from_id > 0 else from_id + ind
|
|
print('%5d %-6s %d' % (ind, 'shortcut', from_id))
|
|
prev_width = out_widths[from_id]
|
|
prev_height = out_heights[from_id]
|
|
prev_filters = out_filters[from_id]
|
|
out_widths.append(prev_width)
|
|
out_heights.append(prev_height)
|
|
out_filters.append(prev_filters)
|
|
elif block['type'] == 'connected':
|
|
filters = int(block['output'])
|
|
print('%5d %-6s %d -> %3d' % (ind, 'connected', prev_filters, filters))
|
|
prev_filters = filters
|
|
out_widths.append(1)
|
|
out_heights.append(1)
|
|
out_filters.append(prev_filters)
|
|
else:
|
|
print('unknown type %s' % (block['type']))
|
|
|
|
|
|
def load_conv(buf, start, conv_model):
|
|
num_w = conv_model.weight.numel()
|
|
num_b = conv_model.bias.numel()
|
|
conv_model.bias.data.copy_(torch.from_numpy(buf[start:start + num_b]));
|
|
start = start + num_b
|
|
conv_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_w]).reshape(conv_model.weight.data.shape));
|
|
start = start + num_w
|
|
return start
|
|
|
|
|
|
def save_conv(fp, conv_model):
|
|
if conv_model.bias.is_cuda:
|
|
convert2cpu(conv_model.bias.data).numpy().tofile(fp)
|
|
convert2cpu(conv_model.weight.data).numpy().tofile(fp)
|
|
else:
|
|
conv_model.bias.data.numpy().tofile(fp)
|
|
conv_model.weight.data.numpy().tofile(fp)
|
|
|
|
|
|
def load_conv_bn(buf, start, conv_model, bn_model):
|
|
num_w = conv_model.weight.numel()
|
|
num_b = bn_model.bias.numel()
|
|
bn_model.bias.data.copy_(torch.from_numpy(buf[start:start + num_b]));
|
|
start = start + num_b
|
|
bn_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_b]));
|
|
start = start + num_b
|
|
bn_model.running_mean.copy_(torch.from_numpy(buf[start:start + num_b]));
|
|
start = start + num_b
|
|
bn_model.running_var.copy_(torch.from_numpy(buf[start:start + num_b]));
|
|
start = start + num_b
|
|
conv_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_w]).reshape(conv_model.weight.data.shape));
|
|
start = start + num_w
|
|
return start
|
|
|
|
|
|
def save_conv_bn(fp, conv_model, bn_model):
|
|
if bn_model.bias.is_cuda:
|
|
convert2cpu(bn_model.bias.data).numpy().tofile(fp)
|
|
convert2cpu(bn_model.weight.data).numpy().tofile(fp)
|
|
convert2cpu(bn_model.running_mean).numpy().tofile(fp)
|
|
convert2cpu(bn_model.running_var).numpy().tofile(fp)
|
|
convert2cpu(conv_model.weight.data).numpy().tofile(fp)
|
|
else:
|
|
bn_model.bias.data.numpy().tofile(fp)
|
|
bn_model.weight.data.numpy().tofile(fp)
|
|
bn_model.running_mean.numpy().tofile(fp)
|
|
bn_model.running_var.numpy().tofile(fp)
|
|
conv_model.weight.data.numpy().tofile(fp)
|
|
|
|
|
|
def load_fc(buf, start, fc_model):
|
|
num_w = fc_model.weight.numel()
|
|
num_b = fc_model.bias.numel()
|
|
fc_model.bias.data.copy_(torch.from_numpy(buf[start:start + num_b]));
|
|
start = start + num_b
|
|
fc_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_w]));
|
|
start = start + num_w
|
|
return start
|
|
|
|
|
|
def save_fc(fp, fc_model):
|
|
fc_model.bias.data.numpy().tofile(fp)
|
|
fc_model.weight.data.numpy().tofile(fp)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import sys
|
|
|
|
blocks = parse_cfg('cfg/yolo.cfg')
|
|
if len(sys.argv) == 2:
|
|
blocks = parse_cfg(sys.argv[1])
|
|
print_cfg(blocks)
|