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

strip_optimizer() default to save FP16 model

This commit is contained in:
Glenn Jocher 2020-06-18 00:16:47 -07:00
parent cce95e744d
commit 10145833c3
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
utils/utils.py
View File

@ -620,6 +620,7 @@ def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_op
# Strip optimizer from *.pt files for lighter files (reduced by 1/2 size) # Strip optimizer from *.pt files for lighter files (reduced by 1/2 size)
x = torch.load(f, map_location=torch.device('cpu')) x = torch.load(f, map_location=torch.device('cpu'))
x['optimizer'] = None x['optimizer'] = None
x['model'].half() # to FP16
torch.save(x, f) torch.save(x, f)
print('Optimizer stripped from %s' % f) print('Optimizer stripped from %s' % f)