34 lines
898 B
Python
34 lines
898 B
Python
import argparse
|
|
import subprocess
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='Process a checkpoint to be published')
|
|
parser.add_argument('in_file', help='input checkpoint filename')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def process_checkpoint(in_file):
|
|
tmp_file = in_file + ".tmp"
|
|
subprocess.Popen(['cp', in_file, tmp_file])
|
|
sha = subprocess.check_output(['sha256sum', tmp_file]).decode()
|
|
out_file = in_file
|
|
if out_file.endswith('.pth'):
|
|
out_file = out_file[:-4]
|
|
final_file = out_file + f'-{sha[:8]}.pth'
|
|
assert final_file != in_file, \
|
|
"The output filename is the same as the input file."
|
|
print("Output file: {}".format(final_file))
|
|
subprocess.Popen(['mv', tmp_file, final_file])
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
process_checkpoint(args.in_file)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|