glenn-jocher
commited on
Commit
•
1014583
1
Parent(s):
cce95e7
strip_optimizer() default to save FP16 model
Browse files- utils/utils.py +1 -0
utils/utils.py
CHANGED
@@ -620,6 +620,7 @@ def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_op
|
|
620 |
# Strip optimizer from *.pt files for lighter files (reduced by 1/2 size)
|
621 |
x = torch.load(f, map_location=torch.device('cpu'))
|
622 |
x['optimizer'] = None
|
|
|
623 |
torch.save(x, f)
|
624 |
print('Optimizer stripped from %s' % f)
|
625 |
|
|
|
620 |
# Strip optimizer from *.pt files for lighter files (reduced by 1/2 size)
|
621 |
x = torch.load(f, map_location=torch.device('cpu'))
|
622 |
x['optimizer'] = None
|
623 |
+
x['model'].half() # to FP16
|
624 |
torch.save(x, f)
|
625 |
print('Optimizer stripped from %s' % f)
|
626 |
|