This commit is contained in:
Glenn Jocher 2019-07-16 17:35:20 +02:00
parent 954deadff3
commit 09b3670579
1 changed files with 4 additions and 7 deletions

View File

@ -167,12 +167,11 @@ def train(cfg,
maps = np.zeros(nc) # mAP per class maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0) # P, R, mAP, F1, test_loss results = (0, 0, 0, 0, 0) # P, R, mAP, F1, test_loss
n_burnin = min(round(nb / 5 + 1), 1000) # burn-in batches n_burnin = min(round(nb / 5 + 1), 1000) # burn-in batches
t, t0 = time.time(), time.time()
torch.cuda.empty_cache() torch.cuda.empty_cache()
for epoch in range(start_epoch, epochs): for epoch in range(start_epoch, epochs):
model.train() model.train()
print(('\n%8s%12s' + '%10s' * 7) % print(('\n%8s' + '%10s' * 8) %
('Epoch', 'Batch', 'GIoU/xy', 'wh', 'obj', 'cls', 'total', 'targets', 'img_size')) ('Epoch', 'GIoU/xy', 'wh', 'obj', 'cls', 'total', 'targets', 'img_size', 'gpu_mem'))
# Update scheduler # Update scheduler
scheduler.step() scheduler.step()
@ -236,10 +235,8 @@ def train(cfg,
# Print batch results # Print batch results
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
# s = ('%8s%12s' + '%10.3g' * 7) % ('%g/%g' % (epoch, epochs - 1), '%g/%g' % (i, nb - 1), *mloss, len(targets), time.time() - t) mem = torch.cuda.memory_cached() if torch.cuda.is_available() else 0
s = ('%8s%12s' + '%10.3g' * 7) % ( s = ('%8s' + '%10.3g' * 8) % ('%g/%g' % (epoch, epochs - 1), *mloss, len(targets), img_size, mem)
'%g/%g' % (epoch, epochs - 1), '%g/%g' % (i, nb - 1), *mloss, len(targets), img_size)
t = time.time()
pbar.set_description(s) # print(s) pbar.set_description(s) # print(s)
# Report time # Report time