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:
1 parent
ca5b10b759
commit
ec1d8496ba
3 files changed
+18
-12
No files matched your search
+2
-2
@@ -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():
|
||||
|
||||
Reference in new issue
Block a user