updates
This commit is contained in:
parent
333cf92bb2
commit
90daf8f19c
41
models.py
41
models.py
|
@ -39,42 +39,37 @@ def create_modules(module_defs):
|
||||||
elif module_def['type'] == 'maxpool':
|
elif module_def['type'] == 'maxpool':
|
||||||
kernel_size = int(module_def['size'])
|
kernel_size = int(module_def['size'])
|
||||||
stride = int(module_def['stride'])
|
stride = int(module_def['stride'])
|
||||||
if kernel_size == 2 and stride == 1:
|
if kernel_size == 2 and stride == 1: # yolov3-tiny
|
||||||
modules.add_module('_debug_padding_%d' % i, nn.ZeroPad2d((0, 1, 0, 1)))
|
modules.add_module('_debug_padding_%d' % i, nn.ZeroPad2d((0, 1, 0, 1)))
|
||||||
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))
|
||||||
modules.add_module('maxpool_%d' % i, maxpool)
|
modules.add_module('maxpool_%d' % i, maxpool)
|
||||||
|
|
||||||
elif module_def['type'] == 'upsample':
|
elif module_def['type'] == 'upsample':
|
||||||
upsample = nn.Upsample(scale_factor=int(module_def['stride']), mode='nearest')
|
modules = nn.Upsample(scale_factor=int(module_def['stride']), mode='nearest')
|
||||||
modules.add_module('upsample_%d' % i, upsample)
|
|
||||||
|
|
||||||
elif module_def['type'] == 'route':
|
elif module_def['type'] == 'route': # nn.Sequential() placeholder for 'route' layer
|
||||||
layers = [int(x) for x in module_def['layers'].split(',')]
|
layers = [int(x) for x in module_def['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])
|
||||||
modules.add_module('route_%d' % i, nn.Sequential()) # Placeholder for 'route' layer
|
|
||||||
# if module_defs[i+1]['type'] == 'reorg3d':
|
# if module_defs[i+1]['type'] == 'reorg3d':
|
||||||
# upsample = nn.Upsample(scale_factor=1/float(module_defs[i+1]['stride']), mode='nearest')
|
# modules = nn.Upsample(scale_factor=1/float(module_defs[i+1]['stride']), mode='nearest') # reorg3d
|
||||||
# modules.add_module('reorg3d_%d' % i, upsample)
|
|
||||||
elif module_def['type'] == 'shortcut':
|
|
||||||
filters = output_filters[int(module_def['from'])]
|
|
||||||
modules.add_module('shortcut_%d' % i, nn.Sequential()) # Placeholder for 'shortcut' layer
|
|
||||||
|
|
||||||
elif module_def['type'] == 'reorg3d':
|
elif module_def['type'] == 'shortcut': # nn.Sequential() placeholder for 'shortcut' layer
|
||||||
|
filters = output_filters[int(module_def['from'])]
|
||||||
|
|
||||||
|
elif module_def['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 module_def['type'] == 'yolo':
|
||||||
yolo_index += 1
|
yolo_index += 1
|
||||||
anchor_idxs = [int(x) for x in module_def['mask'].split(',')]
|
mask = [int(x) for x in module_def['mask'].split(',')] # anchor mask
|
||||||
# Extract anchors
|
a = [float(x) for x in module_def['anchors'].split(',')] # anchors
|
||||||
anchors = [float(x) for x in module_def['anchors'].split(',')]
|
a = [(a[i], a[i + 1]) for i in range(0, len(a), 2)]
|
||||||
anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
|
modules = YOLOLayer(anchors=[a[i] for i in mask], # anchor list
|
||||||
anchors = [anchors[i] for i in anchor_idxs]
|
nc=int(module_def['classes']), # number of classes
|
||||||
nc = int(module_def['classes']) # number of classes
|
img_size=hyperparams['height'], # 416
|
||||||
img_size = hyperparams['height']
|
yolo_index=yolo_index) # 0, 1 or 2
|
||||||
# Define detection layer
|
|
||||||
modules.add_module('yolo_%d' % i, YOLOLayer(anchors, nc, img_size, yolo_index))
|
|
||||||
else:
|
else:
|
||||||
print('Warning: Unrecognized Layer Type: ' + module_def['type'])
|
print('Warning: Unrecognized Layer Type: ' + module_def['type'])
|
||||||
|
|
||||||
|
@ -146,11 +141,11 @@ class YOLOLayer(nn.Module):
|
||||||
io[..., 0:2] = torch.sigmoid(io[..., 0:2]) + self.grid_xy # xy
|
io[..., 0:2] = torch.sigmoid(io[..., 0:2]) + self.grid_xy # xy
|
||||||
io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh # wh yolo method
|
io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh # wh yolo method
|
||||||
io[..., :4] *= self.stride
|
io[..., :4] *= self.stride
|
||||||
|
|
||||||
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
|
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
|
||||||
io[..., 4:] = torch.sigmoid(io[..., 4:]) # p_conf, p_cls
|
io[..., 4:] = torch.sigmoid(io[..., 4:]) # p_conf, p_cls
|
||||||
# io[..., 5:] = F.softmax(io[..., 5:], dim=4) # p_cls
|
# io[..., 5:] = F.softmax(io[..., 5:], dim=4) # p_cls
|
||||||
|
|
||||||
if self.nc == 1:
|
if self.nc == 1:
|
||||||
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235
|
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235
|
||||||
|
|
||||||
|
@ -198,7 +193,7 @@ class Darknet(nn.Module):
|
||||||
layer_i = int(module_def['from'])
|
layer_i = int(module_def['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[0](x, img_size)
|
x = module(x, img_size)
|
||||||
output.append(x)
|
output.append(x)
|
||||||
layer_outputs.append(x)
|
layer_outputs.append(x)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue