From 38804c721b45f92f8139def38e2224a98c66eb0d Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Fri, 8 Oct 2021 17:43:53 -0700
Subject: [PATCH] Checkpoint clean fn useable stand alone

---
 clean_checkpoint.py | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/clean_checkpoint.py b/clean_checkpoint.py
index 34e8604a..3eea15e6 100755
--- a/clean_checkpoint.py
+++ b/clean_checkpoint.py
@@ -20,7 +20,7 @@ parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
                     help='path to latest checkpoint (default: none)')
 parser.add_argument('--output', default='', type=str, metavar='PATH',
                     help='output path')
-parser.add_argument('--use-ema', dest='use_ema', action='store_true',
+parser.add_argument('--no-use-ema', dest='no_use_ema', action='store_true',
                     help='use ema version of weights if present')
 parser.add_argument('--clean-aux-bn', dest='clean_aux_bn', action='store_true',
                     help='remove auxiliary batch norm layers (from SplitBN training) from checkpoint')
@@ -35,19 +35,23 @@ def main():
         print("Error: Output filename ({}) already exists.".format(args.output))
         exit(1)
 
+    clean_checkpoint(args.checkpoint, args.output, not args.no_use_ema, args.clean_aux_bn)
+
+
+def clean_checkpoint(checkpoint, output='', use_ema=True, clean_aux_bn=False):
     # Load an existing checkpoint to CPU, strip everything but the state_dict and re-save
-    if args.checkpoint and os.path.isfile(args.checkpoint):
-        print("=> Loading checkpoint '{}'".format(args.checkpoint))
-        state_dict = load_state_dict(args.checkpoint, use_ema=args.use_ema)
+    if checkpoint and os.path.isfile(checkpoint):
+        print("=> Loading checkpoint '{}'".format(checkpoint))
+        state_dict = load_state_dict(checkpoint, use_ema=use_ema)
         new_state_dict = {}
         for k, v in state_dict.items():
-            if args.clean_aux_bn and 'aux_bn' in k:
+            if clean_aux_bn and 'aux_bn' in k:
                 # If all aux_bn keys are removed, the SplitBN layers will end up as normal and
                 # load with the unmodified model using BatchNorm2d.
                 continue
             name = k[7:] if k.startswith('module') else k
             new_state_dict[name] = v
-        print("=> Loaded state_dict from '{}'".format(args.checkpoint))
+        print("=> Loaded state_dict from '{}'".format(checkpoint))
 
         try:
             torch.save(new_state_dict, _TEMP_NAME, _use_new_zipfile_serialization=False)
@@ -57,17 +61,19 @@ def main():
         with open(_TEMP_NAME, 'rb') as f:
             sha_hash = hashlib.sha256(f.read()).hexdigest()
 
-        if args.output:
-            checkpoint_root, checkpoint_base = os.path.split(args.output)
+        if output:
+            checkpoint_root, checkpoint_base = os.path.split(output)
             checkpoint_base = os.path.splitext(checkpoint_base)[0]
         else:
             checkpoint_root = ''
-            checkpoint_base = os.path.splitext(args.checkpoint)[0]
+            checkpoint_base = os.path.splitext(checkpoint)[0]
         final_filename = '-'.join([checkpoint_base, sha_hash[:8]]) + '.pth'
         shutil.move(_TEMP_NAME, os.path.join(checkpoint_root, final_filename))
         print("=> Saved state_dict to '{}, SHA256: {}'".format(final_filename, sha_hash))
+        return final_filename
     else:
-        print("Error: Checkpoint ({}) doesn't exist".format(args.checkpoint))
+        print("Error: Checkpoint ({}) doesn't exist".format(checkpoint))
+        return ''
 
 
 if __name__ == '__main__':