mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

support bits checkpoints in avg/load

Browse Source
This commit is contained in:
Ross Wightman 2021-10-03 17:31:22 -07:00
parent 007bc39323
commit b2094f4ee8
2 changed files with 15 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
avg_checkpoints.py
View File

@ -41,6 +41,10 @@ def checkpoint_metric(checkpoint_path):
metric = None metric = None
if 'metric' in checkpoint: if 'metric' in checkpoint:
metric = checkpoint['metric'] metric = checkpoint['metric']
elif 'metrics' in checkpoint and 'metric_name' in checkpoint:
metrics = checkpoint['metrics']
print(metrics)
metric = metrics[checkpoint['metric_name']]
return metric return metric

15
timm/models/helpers.py
View File

@ -24,13 +24,20 @@ _logger = logging.getLogger(__name__)
def load_state_dict(checkpoint_path, use_ema=False): def load_state_dict(checkpoint_path, use_ema=False):
if checkpoint_path and os.path.isfile(checkpoint_path): if checkpoint_path and os.path.isfile(checkpoint_path):
checkpoint = torch.load(checkpoint_path, map_location='cpu') checkpoint = torch.load(checkpoint_path, map_location='cpu')
state_dict_key = 'state_dict' state_dict_key = ''
if isinstance(checkpoint, dict): if isinstance(checkpoint, dict):
if use_ema and 'state_dict_ema' in checkpoint: if use_ema and checkpoint.get('state_dict_ema', None) is not None:
state_dict_key = 'state_dict_ema' state_dict_key = 'state_dict_ema'
if state_dict_key and state_dict_key in checkpoint: elif use_ema and checkpoint.get('model_ema', None) is not None:
state_dict_key = 'model_ema'
elif 'state_dict' in checkpoint:
state_dict_key = 'state_dict'
elif 'model' in checkpoint:
state_dict_key = 'model'
if state_dict_key:
state_dict = checkpoint[state_dict_key]
new_state_dict = OrderedDict() new_state_dict = OrderedDict()
for k, v in checkpoint[state_dict_key].items(): for k, v in state_dict.items():
# strip `module.` prefix # strip `module.` prefix
name = k[7:] if k.startswith('module') else k name = k[7:] if k.startswith('module') else k
new_state_dict[name] = v new_state_dict[name] = v