This commit is contained in:
Glenn Jocher 2019-12-20 09:07:25 -08:00
parent 8d54770859
commit 2bc6683325
2 changed files with 40 additions and 24 deletions

45
test.py
View File

@ -78,12 +78,15 @@ def test(cfg,
if batch_i == 0 and not os.path.exists('test_batch0.jpg'): if batch_i == 0 and not os.path.exists('test_batch0.jpg'):
plot_images(imgs=imgs, targets=targets, paths=paths, fname='test_batch0.jpg') plot_images(imgs=imgs, targets=targets, paths=paths, fname='test_batch0.jpg')
# Run model # Disable gradients
inf_out, train_out = model(imgs) # inference and training outputs with torch.no_grad():
# Compute loss # Run model
if hasattr(model, 'hyp'): # if model has loss hyperparameters inf_out, train_out = model(imgs) # inference and training outputs
loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls
# Compute loss
if hasattr(model, 'hyp'): # if model has loss hyperparameters
loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls
# Run NMS # Run NMS
output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres) output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres)
@ -220,12 +223,26 @@ if __name__ == '__main__':
opt = parser.parse_args() opt = parser.parse_args()
print(opt) print(opt)
with torch.no_grad(): # Test
test(opt.cfg, test(opt.cfg,
opt.data, opt.data,
opt.weights, opt.weights,
opt.batch_size, opt.batch_size,
opt.img_size, opt.img_size,
opt.conf_thres, opt.conf_thres,
opt.nms_thres, opt.nms_thres,
opt.save_json or any([x in opt.data for x in ['coco.data', 'coco2014.data', 'coco2017.data']])) opt.save_json or any([x in opt.data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]))
# # Parameter study
# y = []
# x = np.arange(0.4, 0.81, 0.1)
# for v in x:
# y.append(test(opt.cfg, opt.data, opt.weights, opt.batch_size, opt.img_size, 0.1, v, True)[0])
# y = np.stack(y, 0)
#
# fig, ax = plt.subplots(1, 1, figsize=(12, 6))
# ax.plot(x, y[:, 2], marker='.', label='mAP@0.5')
# ax.plot(x, y[:, 3], marker='.', label='mAP@0.5:0.95')
# ax.legend()
# fig.tight_layout()
# plt.savefig('parameters.jpg', dpi=200)

View File

@ -323,16 +323,15 @@ def train():
if opt.prebias: if opt.prebias:
print_model_biases(model) print_model_biases(model)
elif not opt.notest or final_epoch: # Calculate mAP elif not opt.notest or final_epoch: # Calculate mAP
with torch.no_grad(): is_coco = any([x in data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]) and model.nc == 80
is_coco = any([x in data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]) and model.nc == 80 results, maps = test.test(cfg,
results, maps = test.test(cfg, data,
data, batch_size=batch_size,
batch_size=batch_size, img_size=opt.img_size,
img_size=opt.img_size, model=model,
model=model, conf_thres=0.001 if final_epoch else 0.1, # 0.1 for speed
conf_thres=0.001 if final_epoch else 0.1, # 0.1 for speed save_json=final_epoch and is_coco,
save_json=final_epoch and is_coco, dataloader=testloader)
dataloader=testloader)
# Write epoch results # Write epoch results
with open(results_file, 'a') as f: with open(results_file, 'a') as f: