zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

save ckpt in FP16 #119

This commit is contained in:
Glenn Jocher 2020-06-17 22:34:13 -07:00
parent 9fdb0fbacf
commit d9b64c27c2
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
train.py
View File

@ -332,7 +332,7 @@ def train(hyp):
ckpt = {'epoch': epoch, ckpt = {'epoch': epoch,
'best_fitness': best_fitness, 'best_fitness': best_fitness,
'training_results': f.read(), 'training_results': f.read(),
'model': ema.ema.module if hasattr(model, 'module') else ema.ema, 'model': ema.ema.module.half() if hasattr(model, 'module') else ema.ema.half(),
'optimizer': None if final_epoch else optimizer.state_dict()} 'optimizer': None if final_epoch else optimizer.state_dict()}
# Save last, best and delete # Save last, best and delete