updates
This commit is contained in:
parent
d5db50df8e
commit
bce3dd03e8
|
@ -33,7 +33,7 @@ def create_modules(module_defs):
|
||||||
if bn:
|
if bn:
|
||||||
modules.add_module('batch_norm_%d' % i, nn.BatchNorm2d(filters))
|
modules.add_module('batch_norm_%d' % i, nn.BatchNorm2d(filters))
|
||||||
if module_def['activation'] == 'leaky':
|
if module_def['activation'] == 'leaky':
|
||||||
modules.add_module('leaky_%d' % i, nn.LeakyReLU(0.1))
|
modules.add_module('leaky_%d' % i, nn.LeakyReLU(0.1, inplace=True))
|
||||||
|
|
||||||
elif module_def['type'] == 'maxpool':
|
elif module_def['type'] == 'maxpool':
|
||||||
kernel_size = int(module_def['size'])
|
kernel_size = int(module_def['size'])
|
||||||
|
|
3
train.py
3
train.py
|
@ -49,7 +49,6 @@ def train(
|
||||||
start_epoch = 0
|
start_epoch = 0
|
||||||
best_loss = float('inf')
|
best_loss = float('inf')
|
||||||
nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255)
|
nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255)
|
||||||
|
|
||||||
if resume: # Load previously saved model
|
if resume: # Load previously saved model
|
||||||
if transfer: # Transfer learning
|
if transfer: # Transfer learning
|
||||||
chkpt = torch.load(weights + 'yolov3.pt', map_location=device)
|
chkpt = torch.load(weights + 'yolov3.pt', map_location=device)
|
||||||
|
@ -74,7 +73,7 @@ def train(
|
||||||
else:
|
else:
|
||||||
cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
|
cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
|
||||||
|
|
||||||
# Set scheduler (reduce lr at epochs 218, 245, i.e. batches 400k, 450k)
|
# Scheduler (reduce lr at epochs 218, 245, i.e. batches 400k, 450k)
|
||||||
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[218, 245], gamma=0.1,
|
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[218, 245], gamma=0.1,
|
||||||
last_epoch=start_epoch - 1)
|
last_epoch=start_epoch - 1)
|
||||||
|
|
||||||
|
|
|
@ -326,7 +326,6 @@ def build_targets(model, targets):
|
||||||
return txy, twh, tcls, indices
|
return txy, twh, tcls, indices
|
||||||
|
|
||||||
|
|
||||||
# @profile
|
|
||||||
def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
|
def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
|
||||||
"""
|
"""
|
||||||
Removes detections with lower object confidence score than 'conf_thres'
|
Removes detections with lower object confidence score than 'conf_thres'
|
||||||
|
@ -383,6 +382,11 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
|
||||||
dc = pred[pred[:, -1] == c] # select class c
|
dc = pred[pred[:, -1] == c] # select class c
|
||||||
dc = dc[:min(len(dc), 100)] # limit to first 100 boxes: https://github.com/ultralytics/yolov3/issues/117
|
dc = dc[:min(len(dc), 100)] # limit to first 100 boxes: https://github.com/ultralytics/yolov3/issues/117
|
||||||
|
|
||||||
|
# No NMS required if only 1 prediction
|
||||||
|
if len(dc) == 1:
|
||||||
|
det_max.append(dc)
|
||||||
|
continue
|
||||||
|
|
||||||
# Non-maximum suppression
|
# Non-maximum suppression
|
||||||
if nms_style == 'OR': # default
|
if nms_style == 'OR': # default
|
||||||
# METHOD1
|
# METHOD1
|
||||||
|
@ -410,6 +414,9 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
|
||||||
|
|
||||||
elif nms_style == 'MERGE': # weighted mixture box
|
elif nms_style == 'MERGE': # weighted mixture box
|
||||||
while len(dc):
|
while len(dc):
|
||||||
|
if len(dc) == 1:
|
||||||
|
det_max.append(dc)
|
||||||
|
break
|
||||||
i = bbox_iou(dc[0], dc) > nms_thres # iou with other boxes
|
i = bbox_iou(dc[0], dc) > nms_thres # iou with other boxes
|
||||||
weights = dc[i, 4:5]
|
weights = dc[i, 4:5]
|
||||||
dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()
|
dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()
|
||||||
|
|
Loading…
Reference in New Issue