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

Clean checkpoint renames pth w/ SHA hash in a torch hub compatible way

Browse Source
This commit is contained in:
Ross Wightman 2019-11-28 16:42:37 -08:00
parent a39cc43374
commit 3b8f63084d
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
clean_checkpoint.py
View File

@ -2,6 +2,7 @@ import torch
import argparse
import os
import hashlib
import shutil
from collections import OrderedDict
parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
@ -30,11 +31,10 @@ def main():
state_dict_key = 'state_dict_ema' if args.use_ema else 'state_dict'
if state_dict_key in checkpoint:
state_dict = checkpoint[state_dict_key]
else:
print("Error: No state_dict found in checkpoint {}.".format(args.checkpoint))
exit(1)
else:
state_dict = checkpoint
else:
assert False
for k, v in state_dict.items():
name = k[7:] if k.startswith('module') else k
new_state_dict[name] = v
@ -43,7 +43,11 @@ def main():
torch.save(new_state_dict, args.output)
with open(args.output, 'rb') as f:
sha_hash = hashlib.sha256(f.read()).hexdigest()
print("=> Saved state_dict to '{}, SHA256: {}'".format(args.output, sha_hash))
checkpoint_base = os.path.splitext(args.checkpoint)[0]
final_filename = '-'.join([checkpoint_base, sha_hash[:8]]) + '.pth'
shutil.move(args.output, final_filename)
print("=> Saved state_dict to '{}, SHA256: {}'".format(final_filename, sha_hash))
else:
print("Error: Checkpoint ({}) doesn't exist".format(args.checkpoint))