This commit is contained in:
Glenn Jocher 2019-08-03 14:49:38 +02:00
parent 2d8311a83f
commit 658ef8a272
1 changed files with 36 additions and 37 deletions

View File

@ -16,29 +16,29 @@ def create_modules(module_defs):
module_list = nn.ModuleList() module_list = nn.ModuleList()
yolo_index = -1 yolo_index = -1
for i, module_def in enumerate(module_defs): for i, mdef in enumerate(module_defs):
modules = nn.Sequential() modules = nn.Sequential()
if module_def['type'] == 'convolutional': if mdef['type'] == 'convolutional':
bn = int(module_def['batch_normalize']) bn = int(mdef['batch_normalize'])
filters = int(module_def['filters']) filters = int(mdef['filters'])
kernel_size = int(module_def['size']) kernel_size = int(mdef['size'])
pad = (kernel_size - 1) // 2 if int(module_def['pad']) else 0 pad = (kernel_size - 1) // 2 if int(mdef['pad']) else 0
modules.add_module('Conv2d', nn.Conv2d(in_channels=output_filters[-1], modules.add_module('Conv2d', nn.Conv2d(in_channels=output_filters[-1],
out_channels=filters, out_channels=filters,
kernel_size=kernel_size, kernel_size=kernel_size,
stride=int(module_def['stride']), stride=int(mdef['stride']),
padding=pad, padding=pad,
bias=not bn)) bias=not bn))
if bn: if bn:
modules.add_module('BatchNorm2d', nn.BatchNorm2d(filters)) modules.add_module('BatchNorm2d', nn.BatchNorm2d(filters))
if module_def['activation'] == 'leaky': if mdef['activation'] == 'leaky':
# modules.add_module('activation', nn.PReLU(num_parameters=filters, init=0.1)) # modules.add_module('activation', nn.PReLU(num_parameters=filters, init=0.1))
modules.add_module('activation', nn.LeakyReLU(0.1, inplace=True)) modules.add_module('activation', nn.LeakyReLU(0.1, inplace=True))
elif module_def['type'] == 'maxpool': elif mdef['type'] == 'maxpool':
kernel_size = int(module_def['size']) kernel_size = int(mdef['size'])
stride = int(module_def['stride']) stride = int(mdef['stride'])
maxpool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=int((kernel_size - 1) // 2)) maxpool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=int((kernel_size - 1) // 2))
if kernel_size == 2 and stride == 1: # yolov3-tiny if kernel_size == 2 and stride == 1: # yolov3-tiny
modules.add_module('ZeroPad2d', nn.ZeroPad2d((0, 1, 0, 1))) modules.add_module('ZeroPad2d', nn.ZeroPad2d((0, 1, 0, 1)))
@ -46,34 +46,34 @@ def create_modules(module_defs):
else: else:
modules = maxpool modules = maxpool
elif module_def['type'] == 'upsample': elif mdef['type'] == 'upsample':
modules = nn.Upsample(scale_factor=int(module_def['stride']), mode='nearest') modules = nn.Upsample(scale_factor=int(mdef['stride']), mode='nearest')
elif module_def['type'] == 'route': # nn.Sequential() placeholder for 'route' layer elif mdef['type'] == 'route': # nn.Sequential() placeholder for 'route' layer
layers = [int(x) for x in module_def['layers'].split(',')] layers = [int(x) for x in mdef['layers'].split(',')]
filters = sum([output_filters[i + 1 if i > 0 else i] for i in layers]) filters = sum([output_filters[i + 1 if i > 0 else i] for i in layers])
# if module_defs[i+1]['type'] == 'reorg3d': # if mdef[i+1]['type'] == 'reorg3d':
# modules = nn.Upsample(scale_factor=1/float(module_defs[i+1]['stride']), mode='nearest') # reorg3d # modules = nn.Upsample(scale_factor=1/float(mdef[i+1]['stride']), mode='nearest') # reorg3d
elif module_def['type'] == 'shortcut': # nn.Sequential() placeholder for 'shortcut' layer elif mdef['type'] == 'shortcut': # nn.Sequential() placeholder for 'shortcut' layer
filters = output_filters[int(module_def['from'])] filters = output_filters[int(mdef['from'])]
elif module_def['type'] == 'reorg3d': # yolov3-spp-pan-scale elif mdef['type'] == 'reorg3d': # yolov3-spp-pan-scale
# torch.Size([16, 128, 104, 104]) # torch.Size([16, 128, 104, 104])
# torch.Size([16, 64, 208, 208]) <-- # stride 2 interpolate dimensions 2 and 3 to cat with prior layer # torch.Size([16, 64, 208, 208]) <-- # stride 2 interpolate dimensions 2 and 3 to cat with prior layer
pass pass
elif module_def['type'] == 'yolo': elif mdef['type'] == 'yolo':
yolo_index += 1 yolo_index += 1
mask = [int(x) for x in module_def['mask'].split(',')] # anchor mask mask = [int(x) for x in mdef['mask'].split(',')] # anchor mask
a = [float(x) for x in module_def['anchors'].split(',')] # anchors a = [float(x) for x in mdef['anchors'].split(',')] # anchors
a = [(a[i], a[i + 1]) for i in range(0, len(a), 2)] a = [(a[i], a[i + 1]) for i in range(0, len(a), 2)]
modules = YOLOLayer(anchors=[a[i] for i in mask], # anchor list modules = YOLOLayer(anchors=[a[i] for i in mask], # anchor list
nc=int(module_def['classes']), # number of classes nc=int(mdef['classes']), # number of classes
img_size=hyperparams['height'], # 416 img_size=hyperparams['height'], # 416
yolo_index=yolo_index) # 0, 1 or 2 yolo_index=yolo_index) # 0, 1 or 2
else: else:
print('Warning: Unrecognized Layer Type: ' + module_def['type']) print('Warning: Unrecognized Layer Type: ' + mdef['type'])
# Register module list and number of output filters # Register module list and number of output filters
module_list.append(modules) module_list.append(modules)
@ -176,12 +176,12 @@ class Darknet(nn.Module):
layer_outputs = [] layer_outputs = []
output = [] output = []
for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)): for i, (mdef, module) in enumerate(zip(self.module_defs, self.module_list)):
mtype = module_def['type'] mtype = mdef['type']
if mtype in ['convolutional', 'upsample', 'maxpool']: if mtype in ['convolutional', 'upsample', 'maxpool']:
x = module(x) x = module(x)
elif mtype == 'route': elif mtype == 'route':
layer_i = [int(x) for x in module_def['layers'].split(',')] layer_i = [int(x) for x in mdef['layers'].split(',')]
if len(layer_i) == 1: if len(layer_i) == 1:
x = layer_outputs[layer_i[0]] x = layer_outputs[layer_i[0]]
else: else:
@ -192,7 +192,7 @@ class Darknet(nn.Module):
x = torch.cat([layer_outputs[i] for i in layer_i], 1) x = torch.cat([layer_outputs[i] for i in layer_i], 1)
# print(''), [print(layer_outputs[i].shape) for i in layer_i], print(x.shape) # print(''), [print(layer_outputs[i].shape) for i in layer_i], print(x.shape)
elif mtype == 'shortcut': elif mtype == 'shortcut':
layer_i = int(module_def['from']) layer_i = int(mdef['from'])
x = layer_outputs[-1] + layer_outputs[layer_i] x = layer_outputs[-1] + layer_outputs[layer_i]
elif mtype == 'yolo': elif mtype == 'yolo':
x = module(x, img_size) x = module(x, img_size)
@ -226,8 +226,7 @@ class Darknet(nn.Module):
def get_yolo_layers(model): def get_yolo_layers(model):
a = [module_def['type'] == 'yolo' for module_def in model.module_defs] return [i for i, x in enumerate(model.module_defs) if x['type'] == 'yolo'] # [82, 94, 106] for yolov3
return [i for i, x in enumerate(a) if x] # [82, 94, 106] for yolov3
def create_grids(self, img_size=416, ng=(13, 13), device='cpu', type=torch.float32): def create_grids(self, img_size=416, ng=(13, 13), device='cpu', type=torch.float32):
@ -278,10 +277,10 @@ def load_darknet_weights(self, weights, cutoff=-1):
weights = np.fromfile(f, dtype=np.float32) # The rest are weights weights = np.fromfile(f, dtype=np.float32) # The rest are weights
ptr = 0 ptr = 0
for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): for i, (mdef, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
if module_def['type'] == 'convolutional': if mdef['type'] == 'convolutional':
conv_layer = module[0] conv_layer = module[0]
if module_def['batch_normalize']: if mdef['batch_normalize']:
# Load BN bias, weights, running mean and running variance # Load BN bias, weights, running mean and running variance
bn_layer = module[1] bn_layer = module[1]
num_b = bn_layer.bias.numel() # Number of biases num_b = bn_layer.bias.numel() # Number of biases
@ -325,11 +324,11 @@ def save_weights(self, path='model.weights', cutoff=-1):
self.seen.tofile(f) # (int64) number of images seen during training self.seen.tofile(f) # (int64) number of images seen during training
# Iterate through layers # Iterate through layers
for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): for i, (mdef, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
if module_def['type'] == 'convolutional': if mdef['type'] == 'convolutional':
conv_layer = module[0] conv_layer = module[0]
# If batch norm, load bn first # If batch norm, load bn first
if module_def['batch_normalize']: if mdef['batch_normalize']:
bn_layer = module[1] bn_layer = module[1]
bn_layer.bias.data.cpu().numpy().tofile(f) bn_layer.bias.data.cpu().numpy().tofile(f)
bn_layer.weight.data.cpu().numpy().tofile(f) bn_layer.weight.data.cpu().numpy().tofile(f)