timm bits

README.md

@@ -15,6 +15,7 @@
 
 Thanks to the following for hardware support:
 * TPU Research Cloud (TRC) (https://sites.research.google/trc/about/)
+  * TPU support can be found on the [`bits_and_tpu`](https://github.com/rwightman/pytorch-image-models/tree/bits_and_tpu/) branch, w/ some setup help [here](https://github.com/rwightman/pytorch-image-models/tree/bits_and_tpu/timm/bits)
 * Nvidia (https://www.nvidia.com/en-us/)
 
 And a big thanks to all GitHub sponsors who helped with some of my costs before I joined Hugging Face.

clean_checkpoint.py

@@ -13,14 +13,16 @@
 import hashlib
 import shutil
 from collections import OrderedDict
+
 from timm.models.helpers import load_state_dict
+from timm.utils import setup_default_logging
 
 parser = argparse.ArgumentParser(description='PyTorch Checkpoint Cleaner')
 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('--no-use-ema', dest='no_use_ema', action='store_true',
+parser.add_argument('--use-ema', dest='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')
@@ -30,28 +32,25 @@
 
 def main():
     args = parser.parse_args()
+    setup_default_logging()
 
     if os.path.exists(args.output):
         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 checkpoint and os.path.isfile(checkpoint):
-        print("=> Loading checkpoint '{}'".format(checkpoint))
-        state_dict = load_state_dict(checkpoint, use_ema=use_ema)
-        new_state_dict = {}
+    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)
+        new_state_dict = OrderedDict()
         for k, v in state_dict.items():
-            if clean_aux_bn and 'aux_bn' in k:
+            if args.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(checkpoint))
+        print("=> Loaded state_dict from '{}'".format(args.checkpoint))
 
         try:
             torch.save(new_state_dict, _TEMP_NAME, _use_new_zipfile_serialization=False)
@@ -61,19 +60,17 @@ def clean_checkpoint(checkpoint, output='', use_ema=True, clean_aux_bn=False):
         with open(_TEMP_NAME, 'rb') as f:
             sha_hash = hashlib.sha256(f.read()).hexdigest()
 
-        if output:
-            checkpoint_root, checkpoint_base = os.path.split(output)
+        if args.output:
+            checkpoint_root, checkpoint_base = os.path.split(args.output)
             checkpoint_base = os.path.splitext(checkpoint_base)[0]
         else:
             checkpoint_root = ''
-            checkpoint_base = os.path.splitext(checkpoint)[0]
+            checkpoint_base = os.path.splitext(args.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(checkpoint))
-        return ''
+        print("Error: Checkpoint ({}) doesn't exist".format(args.checkpoint))
 
 
 if __name__ == '__main__':

inference.py

@@ -13,7 +13,7 @@
 import torch
 
 from timm.models import create_model, apply_test_time_pool
-from timm.data import ImageDataset, create_loader, resolve_data_config
+from timm.data import ImageDataset, create_loader_v2, resolve_data_config
 from timm.utils import AverageMeter, setup_default_logging
 
 torch.backends.cudnn.benchmark = True
@@ -82,7 +82,7 @@ def main():
     else:
         model = model.cuda()
 
-    loader = create_loader(
+    loader = create_loader_v2(
         ImageDataset(args.data),
         input_size=config['input_size'],
         batch_size=args.batch_size,

launch_xla.py

@@ -0,0 +1,66 @@
+"""
+Adapatation of (pre-elastic) torch.distributed.launch for pytorch xla.
+
+`torch.distributed.launch` is a module that spawns up multiple distributed
+training processes on each of the training nodes.
+
+"""
+
+
+import sys
+import subprocess
+import importlib
+import os
+from argparse import ArgumentParser, REMAINDER
+from typing import Optional, IO
+
+import torch_xla.distributed.xla_multiprocessing as xmp
+
+
+def parse_args():
+    """
+    Helper function parsing the command line options
+    @retval ArgumentParser
+    """
+    parser = ArgumentParser(
+        description="PyTorch distributed training launch helper utility"
+                    "that will spawn up multiple distributed processes")
+
+    # Optional arguments for the launch helper
+    parser.add_argument("--num-devices", type=int, default=1,
+                        help="The number of XLA devices to use for distributed training")
+
+    # positional
+    parser.add_argument(
+        "script", type=str,
+        help="The full path to the single device training script to be launched"
+             "in parallel, followed by all the arguments for the training script")
+
+    # rest from the training program
+    parser.add_argument('script_args', nargs=REMAINDER)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # set PyTorch distributed related environmental variables
+    # current_env = os.environ.copy()
+    # current_env["MASTER_ADDR"] = args.master_addr
+    # current_env["MASTER_PORT"] = str(args.master_port)
+    # current_env["WORLD_SIZE"] = str(dist_world_size)
+    # if 'OMP_NUM_THREADS' not in os.environ and args.nproc_per_node > 1:
+    #    current_env["OMP_NUM_THREADS"] = str(1)
+
+    script_abs = os.path.abspath(args.script)
+    script_base, script_rel = os.path.split(script_abs)
+    sys.path.append(script_base)
+    mod = importlib.import_module(os.path.splitext(script_rel)[0])
+
+    sys.argv = [args.script] + args.script_args
+
+    xmp.spawn(mod._mp_entry, args=(), nprocs=args.num_devices)
+
+
+if __name__ == "__main__":
+    main()