|
|
|
@ -268,8 +268,23 @@ class ModelEMA:
|
|
|
|
|
copy_attr(self.ema, model, include, exclude)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer()
|
|
|
|
|
# Strip optimizer from 'f' to finalize training, optionally save as 's'
|
|
|
|
|
def strip_optimizer(f='best.pt', s=''):
|
|
|
|
|
"""
|
|
|
|
|
Strip optimizer from 'f' to finalize training, optionally save as 's'.
|
|
|
|
|
|
|
|
|
|
Usage:
|
|
|
|
|
from ultralytics.yolo.utils.torch_utils import strip_optimizer
|
|
|
|
|
from pathlib import Path
|
|
|
|
|
for f in Path('/Users/glennjocher/Downloads/weights').glob('*.pt'):
|
|
|
|
|
strip_optimizer(f)
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
f (str): file path to model state to strip the optimizer from. Default is 'best.pt'.
|
|
|
|
|
s (str): file path to save the model with stripped optimizer to. Default is ''. If not provided, the original file will be overwritten.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
"""
|
|
|
|
|
x = torch.load(f, map_location=torch.device('cpu'))
|
|
|
|
|
args = {**DEFAULT_CONFIG_DICT, **x['train_args']} # combine model args with default args, preferring model args
|
|
|
|
|
if x.get('ema'):
|
|
|
|
|