Improved model+EMA checkpointing (#2292)

* Enhanced model+EMA checkpointing

* update

* bug fix

* bug fix 2

* always save optimizer

* ema half

* remove model.float()

* model half

* carry ema/model in fp32

* rm model.float()

* both to float always

* cleanup

* cleanup
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-02-24 18:26:46 -08:00
1 parent ca5b10b759
commit ec1d8496ba
3 files changed
+18 -12

No files matched your search

+2 -2
View File
@@ -484,8 +484,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
def strip_optimizer(f='weights/best.pt', s=''): # from utils.general import *; strip_optimizer()
# Strip optimizer from 'f' to finalize training, optionally save as 's'
x = torch.load(f, map_location=torch.device('cpu'))
for key in 'optimizer', 'training_results', 'wandb_id':
x[key] = None
for k in 'optimizer', 'training_results', 'wandb_id', 'ema': # keys
x[k] = None
x['epoch'] = -1
x['model'].half() # to FP16
for p in x['model'].parameters():