File size: 2,252 Bytes
789eef1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import argparse
import os
import torch
from mmengine.runner.checkpoint import _load_checkpoint
def parse_args():
parser = argparse.ArgumentParser(description='Clean checkpoints by removing unnecessary data')
parser.add_argument('--checkpoint_dir', help='Directory containing checkpoints')
parser.add_argument('--checkpoint_path', help='Path to a single checkpoint')
args = parser.parse_args()
if not args.checkpoint_dir and not args.checkpoint_path:
print("\033[91m" + "Error: Either checkpoint_dir or checkpoint_path must be provided." + "\033[0m")
exit(1)
if args.checkpoint_path and not os.path.exists(args.checkpoint_path):
print("\033[91m" + "Error: Checkpoint path does not exist." + "\033[0m")
exit(1)
if args.checkpoint_dir and not os.path.exists(args.checkpoint_dir):
print("\033[91m" + "Error: Checkpoint directory does not exist." + "\033[0m")
exit(1)
return args
def clean_checkpoint(checkpoint_path):
print("\033[96m" + f"Cleaning checkpoint: {checkpoint_path}" + "\033[0m")
checkpoint = _load_checkpoint(checkpoint_path)
# Remove unnecessary parts of the checkpoint
clean_checkpoint = {'state_dict': checkpoint['state_dict'], 'meta': checkpoint['meta']}
# new_checkpoint_path = os.path.join(os.path.dirname(checkpoint_path), os.path.basename(checkpoint_path).replace('.pth', '_clean.pth'))
new_checkpoint_path = checkpoint_path
print('\033[93m' + f'Saving cleaned checkpoint to {new_checkpoint_path}' + '\033[0m')
torch.save(clean_checkpoint, new_checkpoint_path)
return
def main():
args = parse_args()
if args.checkpoint_dir:
for root, _, files in os.walk(args.checkpoint_dir):
for file in files:
if file.endswith('.pth'):
checkpoint_path = os.path.join(root, file)
clean_checkpoint(checkpoint_path)
elif args.checkpoint_path:
clean_checkpoint(args.checkpoint_path)
if __name__ == "__main__":
main()
|