Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Update train.py (#2290)

Browse Source 
* Update train.py

* Update train.py

* Update train.py

* Update train.py

* Create train.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-02-24 13:31:20 -08:00
1 parent 0070995bd5
commit ca5b10b759
1 file changed
+16 -19
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+16 -19
View File
@@ -146,8 +146,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Results # Results
if ckpt.get('training_results') is not None: if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file: results_file.write_text(ckpt['training_results']) # write results.txt
file.write(ckpt['training_results']) # write results.txt
# Epochs # Epochs
start_epoch = ckpt['epoch'] + 1 start_epoch = ckpt['epoch'] + 1
@@ -354,7 +353,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Write # Write
with open(results_file, 'a') as f: with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results + '\n') # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss
if len(opt.name) and opt.bucket: if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name)) os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
@@ -375,15 +374,13 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
best_fitness = fi best_fitness = fi
# Save model # Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve) if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
if save: ckpt = {'epoch': epoch,
with open(results_file, 'r') as f: # create checkpoint 'best_fitness': best_fitness,
ckpt = {'epoch': epoch, 'training_results': results_file.read_text(),
'best_fitness': best_fitness, 'model': ema.ema,
'training_results': f.read(), 'optimizer': None if final_epoch else optimizer.state_dict(),
'model': ema.ema, 'wandb_id': wandb_run.id if wandb else None}
'optimizer': None if final_epoch else optimizer.state_dict(),
'wandb_id': wandb_run.id if wandb else None}
# Save last, best and delete # Save last, best and delete
torch.save(ckpt, last) torch.save(ckpt, last)
@@ -396,9 +393,9 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
if rank in [-1, 0]: if rank in [-1, 0]:
# Strip optimizers # Strip optimizers
final = best if best.exists() else last # final model final = best if best.exists() else last # final model
for f in [last, best]: for f in last, best:
if f.exists(): if f.exists():
strip_optimizer(f) # strip optimizers strip_optimizer(f)
if opt.bucket: if opt.bucket:
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
@@ -415,17 +412,17 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Test best.pt # Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600)) logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for conf, iou, save_json in ([0.25, 0.45, False], [0.001, 0.65, True]): # speed, mAP tests for m in (last, best) if best.exists() else (last): # speed, mAP tests
results, _, _ = test.test(opt.data, results, _, _ = test.test(opt.data,
batch_size=batch_size * 2, batch_size=batch_size * 2,
imgsz=imgsz_test, imgsz=imgsz_test,
conf_thres=conf, conf_thres=0.001,
iou_thres=iou, iou_thres=0.7,
model=attempt_load(final, device).half(), model=attempt_load(m, device).half(),
single_cls=opt.single_cls, single_cls=opt.single_cls,
dataloader=testloader, dataloader=testloader,
save_dir=save_dir, save_dir=save_dir,
save_json=save_json, save_json=True,
plots=False) plots=False)
else: else: