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()