HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.0.174 Tuner plots and improvements (#4799)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-09-10 03:27:23 +02:00 • committed by GitHub
parent dfe6dfb1d2
commit 16ce193d6e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 248 additions and 88 deletions
Download patch file Download diff file Expand all files Collapse all files

10
ultralytics/utils/torch_utils.py
View file

@ -395,12 +395,6 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None:
strip_optimizer(f)
```
"""
# Use dill (if exists) to serialize the lambda functions where pickle does not do this
try:
import dill as pickle
except ImportError:
import pickle
x = torch.load(f, map_location=torch.device('cpu'))
if 'model' not in x:
LOGGER.info(f'Skipping {f}, not a valid Ultralytics model.')
@ -419,8 +413,8 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None:
p.requires_grad = False
x['train_args'] = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS} # strip non-default keys
# x['model'].args = x['train_args']
torch.save(x, s or f, pickle_module=pickle)
mb = os.path.getsize(s or f) / 1E6 # filesize
torch.save(x, s or f)
mb = os.path.getsize(s or f) / 1E6 # file size
LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")