updates
This commit is contained in:
parent
2d8311a83f
commit
658ef8a272
73
models.py
73
models.py
|
@ -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)
|
||||||
|
|
Loading…
Reference in New Issue