updates
This commit is contained in:
parent
62761cffe6
commit
22dc8c0ea6
64
models.py
64
models.py
|
@ -141,7 +141,7 @@ class YOLOLayer(nn.Module):
|
|||
self.grid_xy = torch.cat((self.grid_x, self.grid_y), 2)
|
||||
self.anchor_wh = torch.cat((self.anchor_w, self.anchor_h), 2) / nG
|
||||
|
||||
def forward(self, p, targets=None, batch_report=False, var=None):
|
||||
def forward(self, p, targets=None, var=None):
|
||||
FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor
|
||||
bs = 1 if ONNX_EXPORT else p.shape[0] # batch size
|
||||
nG = self.nG # number of grid points
|
||||
|
@ -178,18 +178,7 @@ class YOLOLayer(nn.Module):
|
|||
# width = ((w.data * 2) ** 2) * self.anchor_w
|
||||
# height = ((h.data * 2) ** 2) * self.anchor_h
|
||||
|
||||
p_boxes = None
|
||||
if batch_report:
|
||||
# Predicted boxes: add offset and scale with anchors (in grid space, i.e. 0-13)
|
||||
gx = x.data + self.grid_x[:, :, :nG, :nG]
|
||||
gy = y.data + self.grid_y[:, :, :nG, :nG]
|
||||
p_boxes = torch.stack((gx - width / 2,
|
||||
gy - height / 2,
|
||||
gx + width / 2,
|
||||
gy + height / 2), 4) # x1y1x2y2
|
||||
|
||||
tx, ty, tw, th, mask, tcls, TP, FP, FN, TC = \
|
||||
build_targets(p_boxes, p_conf, p_cls, targets, self.anchor_wh, self.nA, self.nC, nG, batch_report)
|
||||
tx, ty, tw, th, mask, tcls = build_targets(targets, self.anchor_wh, self.nA, self.nC, nG)
|
||||
|
||||
tcls = tcls[mask]
|
||||
if x.is_cuda:
|
||||
|
@ -214,26 +203,9 @@ class YOLOLayer(nn.Module):
|
|||
lconf = (k * 64) * BCEWithLogitsLoss(p_conf, mask.float())
|
||||
|
||||
# Sum loss components
|
||||
balance_losses_flag = False
|
||||
if balance_losses_flag:
|
||||
k = 1 / self.loss_means.clone()
|
||||
loss = (lx * k[0] + ly * k[1] + lw * k[2] + lh * k[3] + lconf * k[4] + lcls * k[5]) / k.mean()
|
||||
loss = lx + ly + lw + lh + lconf + lcls
|
||||
|
||||
self.loss_means = self.loss_means * 0.99 + \
|
||||
FT([lx.data, ly.data, lw.data, lh.data, lconf.data, lcls.data]) * 0.01
|
||||
else:
|
||||
loss = lx + ly + lw + lh + lconf + lcls
|
||||
|
||||
# Sum False Positives from unassigned anchors
|
||||
FPe = torch.zeros(self.nC)
|
||||
if batch_report:
|
||||
i = torch.sigmoid(p_conf[~mask]) > 0.5
|
||||
if i.sum() > 0:
|
||||
FP_classes = torch.argmax(p_cls[~mask][i], 1)
|
||||
FPe = torch.bincount(FP_classes, minlength=self.nC).float().cpu() # extra FPs
|
||||
|
||||
return loss, loss.item(), lx.item(), ly.item(), lw.item(), lh.item(), lconf.item(), lcls.item(), \
|
||||
nT, TP, FP, FPe, FN, TC
|
||||
return loss, loss.item(), lx.item(), ly.item(), lw.item(), lh.item(), lconf.item(), lcls.item(), nT
|
||||
|
||||
else:
|
||||
if ONNX_EXPORT:
|
||||
|
@ -273,9 +245,9 @@ class Darknet(nn.Module):
|
|||
self.module_defs[0]['height'] = img_size
|
||||
self.hyperparams, self.module_list = create_modules(self.module_defs)
|
||||
self.img_size = img_size
|
||||
self.loss_names = ['loss', 'x', 'y', 'w', 'h', 'conf', 'cls', 'nT', 'TP', 'FP', 'FPe', 'FN', 'TC']
|
||||
self.loss_names = ['loss', 'x', 'y', 'w', 'h', 'conf', 'cls', 'nT']
|
||||
|
||||
def forward(self, x, targets=None, batch_report=False, var=0):
|
||||
def forward(self, x, targets=None, var=0):
|
||||
self.losses = defaultdict(float)
|
||||
is_training = targets is not None
|
||||
layer_outputs = []
|
||||
|
@ -296,7 +268,7 @@ class Darknet(nn.Module):
|
|||
elif module_def['type'] == 'yolo':
|
||||
# Train phase: get loss
|
||||
if is_training:
|
||||
x, *losses = module[0](x, targets, batch_report, var)
|
||||
x, *losses = module[0](x, targets, var)
|
||||
for name, loss in zip(self.loss_names, losses):
|
||||
self.losses[name] += loss
|
||||
# Test phase: Get detections
|
||||
|
@ -306,29 +278,7 @@ class Darknet(nn.Module):
|
|||
layer_outputs.append(x)
|
||||
|
||||
if is_training:
|
||||
if batch_report:
|
||||
self.losses['TC'] /= 3 # target category
|
||||
metrics = torch.zeros(3, len(self.losses['FPe'])) # TP, FP, FN
|
||||
|
||||
ui = np.unique(self.losses['TC'])[1:]
|
||||
for i in ui:
|
||||
j = self.losses['TC'] == float(i)
|
||||
metrics[0, i] = (self.losses['TP'][j] > 0).sum().float() # TP
|
||||
metrics[1, i] = (self.losses['FP'][j] > 0).sum().float() # FP
|
||||
metrics[2, i] = (self.losses['FN'][j] == 3).sum().float() # FN
|
||||
metrics[1] += self.losses['FPe']
|
||||
|
||||
self.losses['TP'] = metrics[0].sum()
|
||||
self.losses['FP'] = metrics[1].sum()
|
||||
self.losses['FN'] = metrics[2].sum()
|
||||
self.losses['metrics'] = metrics
|
||||
else:
|
||||
self.losses['TP'] = 0
|
||||
self.losses['FP'] = 0
|
||||
self.losses['FN'] = 0
|
||||
|
||||
self.losses['nT'] /= 3
|
||||
self.losses['TC'] = 0
|
||||
|
||||
if ONNX_EXPORT:
|
||||
output = torch.cat(output, 1) # merge the 3 layers 85 x (507, 2028, 8112) to 85 x 10647
|
||||
|
|
36
train.py
36
train.py
|
@ -20,7 +20,6 @@ def train(
|
|||
batch_size=16,
|
||||
accumulated_batches=1,
|
||||
weights='weights',
|
||||
report=False,
|
||||
multi_scale=False,
|
||||
freeze_backbone=True,
|
||||
var=0,
|
||||
|
@ -30,7 +29,7 @@ def train(
|
|||
if multi_scale: # pass maximum multi_scale size
|
||||
img_size = 608
|
||||
else:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
torch.backends.cudnn.benchmark = True # unsuitable for multiscale
|
||||
|
||||
latest = os.path.join(weights, 'latest.pt')
|
||||
best = os.path.join(weights, 'best.pt')
|
||||
|
@ -93,12 +92,11 @@ def train(
|
|||
|
||||
model_info(model)
|
||||
t0 = time.time()
|
||||
mean_recall, mean_precision = 0, 0
|
||||
for epoch in range(epochs):
|
||||
epoch += start_epoch
|
||||
|
||||
print(('%8s%12s' + '%10s' * 14) % ('Epoch', 'Batch', 'x', 'y', 'w', 'h', 'conf', 'cls', 'total', 'P', 'R',
|
||||
'nTargets', 'TP', 'FP', 'FN', 'time'))
|
||||
print(('%8s%12s' + '%10s' * 9) % (
|
||||
'Epoch', 'Batch', 'x', 'y', 'w', 'h', 'conf', 'cls', 'total', 'nTargets', 'time'))
|
||||
|
||||
# Update scheduler (automatic)
|
||||
# scheduler.step()
|
||||
|
@ -124,7 +122,6 @@ def train(
|
|||
|
||||
ui = -1
|
||||
rloss = defaultdict(float) # running loss
|
||||
metrics = torch.zeros(3, num_classes)
|
||||
optimizer.zero_grad()
|
||||
for i, (imgs, targets) in enumerate(dataloader):
|
||||
if sum([len(x) for x in targets]) < 1: # if no targets continue
|
||||
|
@ -137,7 +134,7 @@ def train(
|
|||
g['lr'] = lr
|
||||
|
||||
# Compute loss, compute gradient, update parameters
|
||||
loss = model(imgs.to(device), targets, batch_report=report, var=var)
|
||||
loss = model(imgs.to(device), targets, var=var)
|
||||
loss.backward()
|
||||
|
||||
# accumulate gradient for x batches before optimizing
|
||||
|
@ -150,27 +147,10 @@ def train(
|
|||
for key, val in model.losses.items():
|
||||
rloss[key] = (rloss[key] * ui + val) / (ui + 1)
|
||||
|
||||
if report:
|
||||
TP, FP, FN = metrics
|
||||
metrics += model.losses['metrics']
|
||||
|
||||
# Precision
|
||||
precision = TP / (TP + FP)
|
||||
k = (TP + FP) > 0
|
||||
if k.sum() > 0:
|
||||
mean_precision = precision[k].mean()
|
||||
|
||||
# Recall
|
||||
recall = TP / (TP + FN)
|
||||
k = (TP + FN) > 0
|
||||
if k.sum() > 0:
|
||||
mean_recall = recall[k].mean()
|
||||
|
||||
s = ('%8s%12s' + '%10.3g' * 14) % (
|
||||
s = ('%8s%12s' + '%10.3g' * 9) % (
|
||||
'%g/%g' % (epoch, epochs - 1), '%g/%g' % (i, len(dataloader) - 1), rloss['x'],
|
||||
rloss['y'], rloss['w'], rloss['h'], rloss['conf'], rloss['cls'],
|
||||
rloss['loss'], mean_precision, mean_recall, model.losses['nT'], model.losses['TP'],
|
||||
model.losses['FP'], model.losses['FN'], time.time() - t0)
|
||||
rloss['loss'], model.losses['nT'], time.time() - t0)
|
||||
t0 = time.time()
|
||||
print(s)
|
||||
|
||||
|
@ -214,9 +194,8 @@ if __name__ == '__main__':
|
|||
parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
|
||||
parser.add_argument('--weights', type=str, default='weights', help='path to store weights')
|
||||
parser.add_argument('--resume', action='store_true', help='resume training flag')
|
||||
parser.add_argument('--report', action='store_true', help='report TP, FP, FN, P and R per batch (slower)')
|
||||
parser.add_argument('--freeze', action='store_true', help='freeze darknet53.conv.74 layers for first epoch')
|
||||
parser.add_argument('--var', type=float, default=0, help='optional test variable')
|
||||
parser.add_argument('--var', type=float, default=0, help='test variable')
|
||||
opt = parser.parse_args()
|
||||
print(opt, end='\n\n')
|
||||
|
||||
|
@ -231,7 +210,6 @@ if __name__ == '__main__':
|
|||
batch_size=opt.batch_size,
|
||||
accumulated_batches=opt.accumulated_batches,
|
||||
weights=opt.weights,
|
||||
report=opt.report,
|
||||
multi_scale=opt.multi_scale,
|
||||
freeze_backbone=opt.freeze,
|
||||
var=opt.var,
|
||||
|
|
|
@ -214,7 +214,7 @@ def bbox_iou(box1, box2, x1y1x2y2=True):
|
|||
return inter_area / (b1_area + b2_area - inter_area + 1e-16)
|
||||
|
||||
|
||||
def build_targets(pred_boxes, pred_conf, pred_cls, target, anchor_wh, nA, nC, nG, batch_report):
|
||||
def build_targets(target, anchor_wh, nA, nC, nG):
|
||||
"""
|
||||
returns nT, nCorrect, tx, ty, tw, th, tconf, tcls
|
||||
"""
|
||||
|
@ -226,9 +226,6 @@ def build_targets(pred_boxes, pred_conf, pred_cls, target, anchor_wh, nA, nC, nG
|
|||
th = torch.zeros(nB, nA, nG, nG)
|
||||
tconf = torch.ByteTensor(nB, nA, nG, nG).fill_(0)
|
||||
tcls = torch.ByteTensor(nB, nA, nG, nG, nC).fill_(0) # nC = number of classes
|
||||
TP = torch.ByteTensor(nB, max(nT)).fill_(0)
|
||||
FP = torch.ByteTensor(nB, max(nT)).fill_(0)
|
||||
FN = torch.ByteTensor(nB, max(nT)).fill_(0)
|
||||
TC = torch.ShortTensor(nB, max(nT)).fill_(-1) # target category
|
||||
|
||||
for b in range(nB):
|
||||
|
@ -293,18 +290,7 @@ def build_targets(pred_boxes, pred_conf, pred_cls, target, anchor_wh, nA, nC, nG
|
|||
tcls[b, a, gj, gi, tc] = 1
|
||||
tconf[b, a, gj, gi] = 1
|
||||
|
||||
if batch_report:
|
||||
# predicted classes and confidence
|
||||
tb = torch.cat((gx - gw / 2, gy - gh / 2, gx + gw / 2, gy + gh / 2)).view(4, -1).t() # target boxes
|
||||
pcls = torch.argmax(pred_cls[b, a, gj, gi], 1).cpu()
|
||||
pconf = torch.sigmoid(pred_conf[b, a, gj, gi]).cpu()
|
||||
iou_pred = bbox_iou(tb, pred_boxes[b, a, gj, gi].cpu())
|
||||
|
||||
TP[b, i] = (pconf > 0.5) & (iou_pred > 0.5) & (pcls == tc)
|
||||
FP[b, i] = (pconf > 0.5) & (TP[b, i] == 0) # coordinates or class are wrong
|
||||
FN[b, i] = pconf <= 0.5 # confidence score is too low (set to zero)
|
||||
|
||||
return tx, ty, tw, th, tconf, tcls, TP, FP, FN, TC
|
||||
return tx, ty, tw, th, tconf, tcls
|
||||
|
||||
|
||||
def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
|
||||
|
|
Loading…
Reference in New Issue