Update README.md

* Add support for TF weights and modelling specifics to MaxVit (testing ported weights)
* More fine-tuned CLIP ViT configs
* ConvNeXt and MaxVit updated to new pretrained cfgs use
* EfficientNetV2, MaxVit and ConvNeXt high res models use squash crop/resize

README.md

@@ -21,6 +21,68 @@ And a big thanks to all GitHub sponsors who helped with some of my costs before
 
 ## What's New
 
+# Dec 5, 2022
+
+* Pre-release (`0.8.0dev0`) of multi-weight support (`model_arch.pretrained_tag`). Install with `pip install --pre timm`
+  * vision_transformer, maxvit, convnext are the first three model impl w/ support
+  * model names are changing with this (previous _21k, etc. fn will merge), still sorting out deprecation handling
+  * bugs are likely, but I need feedback so please try it out
+  * if stability is needed, please use 0.6.x pypi releases or clone from [0.6.x branch](https://github.com/rwightman/pytorch-image-models/tree/0.6.x)
+* Support for PyTorch 2.0 compile is added in train/validate/inference/benchmark, use `--torchcompile` argument
+* Inference script allows more control over output, select k for top-class index + prob json, csv or parquet output
+* Add a full set of fine-tuned CLIP image tower weights from both LAION-2B and original OpenAI CLIP models
+
+| model                                            |   top1 |   param_count |   gmac |   macts | hub                                                                                  |
+|:-------------------------------------------------|-------:|--------------:|-------:|--------:|:-------------------------------------------------------------------------------------|
+| vit_huge_patch14_clip_336.laion2b_ft_in12k_in1k  |   88.6 |         632.5 |  391   |   407.5 | [link](https://huggingface.co/timm/vit_huge_patch14_clip_336.laion2b_ft_in12k_in1k)  |
+| vit_large_patch14_clip_336.openai_ft_in12k_in1k  |   88.3 |         304.5 |  191.1 |   270.2 | [link](https://huggingface.co/timm/vit_large_patch14_clip_336.openai_ft_in12k_in1k)  |
+| vit_huge_patch14_clip_224.laion2b_ft_in12k_in1k  |   88.2 |         632   |  167.4 |   139.4 | [link](https://huggingface.co/timm/vit_huge_patch14_clip_224.laion2b_ft_in12k_in1k)  |
+| vit_large_patch14_clip_336.laion2b_ft_in12k_in1k |   88.2 |         304.5 |  191.1 |   270.2 | [link](https://huggingface.co/timm/vit_large_patch14_clip_336.laion2b_ft_in12k_in1k) |
+| vit_large_patch14_clip_224.openai_ft_in12k_in1k  |   88.2 |         304.2 |   81.1 |    88.8 | [link](https://huggingface.co/timm/vit_large_patch14_clip_224.openai_ft_in12k_in1k)  |
+| vit_large_patch14_clip_224.laion2b_ft_in12k_in1k |   87.9 |         304.2 |   81.1 |    88.8 | [link](https://huggingface.co/timm/vit_large_patch14_clip_224.laion2b_ft_in12k_in1k) |
+| vit_large_patch14_clip_224.openai_ft_in1k        |   87.9 |         304.2 |   81.1 |    88.8 | [link](https://huggingface.co/timm/vit_large_patch14_clip_224.openai_ft_in1k)        |
+| vit_large_patch14_clip_336.laion2b_ft_in1k       |   87.9 |         304.5 |  191.1 |   270.2 | [link](https://huggingface.co/timm/vit_large_patch14_clip_336.laion2b_ft_in1k)       |
+| vit_huge_patch14_clip_224.laion2b_ft_in1k        |   87.6 |         632   |  167.4 |   139.4 | [link](https://huggingface.co/timm/vit_huge_patch14_clip_224.laion2b_ft_in1k)        |
+| vit_large_patch14_clip_224.laion2b_ft_in1k       |   87.3 |         304.2 |   81.1 |    88.8 | [link](https://huggingface.co/timm/vit_large_patch14_clip_224.laion2b_ft_in1k)       |
+| vit_base_patch16_clip_384.laion2b_ft_in12k_in1k  |   87.2 |          86.9 |   55.5 |   101.6 | [link](https://huggingface.co/timm/vit_base_patch16_clip_384.laion2b_ft_in12k_in1k)  |
+| vit_base_patch16_clip_384.openai_ft_in12k_in1k   |   87   |          86.9 |   55.5 |   101.6 | [link](https://huggingface.co/timm/vit_base_patch16_clip_384.openai_ft_in12k_in1k)   |
+| vit_base_patch16_clip_384.laion2b_ft_in1k        |   86.6 |          86.9 |   55.5 |   101.6 | [link](https://huggingface.co/timm/vit_base_patch16_clip_384.laion2b_ft_in1k)        |
+| vit_base_patch16_clip_384.openai_ft_in1k         |   86.2 |          86.9 |   55.5 |   101.6 | [link](https://huggingface.co/timm/vit_base_patch16_clip_384.openai_ft_in1k)         |
+| vit_base_patch16_clip_224.laion2b_ft_in12k_in1k  |   86.2 |          86.6 |   17.6 |    23.9 | [link](https://huggingface.co/timm/vit_base_patch16_clip_224.laion2b_ft_in12k_in1k)  |
+| vit_base_patch16_clip_224.openai_ft_in12k_in1k   |   85.9 |          86.6 |   17.6 |    23.9 | [link](https://huggingface.co/timm/vit_base_patch16_clip_224.openai_ft_in12k_in1k)   |
+| vit_base_patch32_clip_448.laion2b_ft_in12k_in1k  |   85.8 |          88.3 |   17.9 |    23.9 | [link](https://huggingface.co/timm/vit_base_patch32_clip_448.laion2b_ft_in12k_in1k)  |
+| vit_base_patch16_clip_224.laion2b_ft_in1k        |   85.5 |          86.6 |   17.6 |    23.9 | [link](https://huggingface.co/timm/vit_base_patch16_clip_224.laion2b_ft_in1k)        |
+| vit_base_patch32_clip_384.laion2b_ft_in12k_in1k  |   85.4 |          88.3 |   13.1 |    16.5 | [link](https://huggingface.co/timm/vit_base_patch32_clip_384.laion2b_ft_in12k_in1k)  |
+| vit_base_patch16_clip_224.openai_ft_in1k         |   85.3 |          86.6 |   17.6 |    23.9 | [link](https://huggingface.co/timm/vit_base_patch16_clip_224.openai_ft_in1k)         |
+| vit_base_patch32_clip_384.openai_ft_in12k_in1k   |   85.2 |          88.3 |   13.1 |    16.5 | [link](https://huggingface.co/timm/vit_base_patch32_clip_384.openai_ft_in12k_in1k)   |
+| vit_base_patch32_clip_224.laion2b_ft_in12k_in1k  |   83.3 |          88.2 |    4.4 |     5   | [link](https://huggingface.co/timm/vit_base_patch32_clip_224.laion2b_ft_in12k_in1k)  |
+| vit_base_patch32_clip_224.laion2b_ft_in1k        |   82.6 |          88.2 |    4.4 |     5   | [link](https://huggingface.co/timm/vit_base_patch32_clip_224.laion2b_ft_in1k)        |
+| vit_base_patch32_clip_224.openai_ft_in1k         |   81.9 |          88.2 |    4.4 |     5   | [link](https://huggingface.co/timm/vit_base_patch32_clip_224.openai_ft_in1k)         |
+
+* Port of MaxViT Tensorflow Weights from official impl at https://github.com/google-research/maxvit
+  * There was larger than expected drops for the upscaled 384/512 in21k fine-tune weights, possible detail missing, but the 21k FT did seem sensitive to small preprocessing
+
+| model                              |   top1 |   param_count |   gmac |   macts | hub                                                                    |
+|:-----------------------------------|-------:|--------------:|-------:|--------:|:-----------------------------------------------------------------------|
+| maxvit_xlarge_tf_512.in21k_ft_in1k |   88.5 |         475.8 |  534.1 |  1413.2 | [link](https://huggingface.co/timm/maxvit_xlarge_tf_512.in21k_ft_in1k) |
+| maxvit_xlarge_tf_384.in21k_ft_in1k |   88.3 |         475.3 |  292.8 |   668.8 | [link](https://huggingface.co/timm/maxvit_xlarge_tf_384.in21k_ft_in1k) |
+| maxvit_base_tf_512.in21k_ft_in1k   |   88.2 |         119.9 |  138   |   704   | [link](https://huggingface.co/timm/maxvit_base_tf_512.in21k_ft_in1k)   |
+| maxvit_large_tf_512.in21k_ft_in1k  |   88   |         212.3 |  244.8 |   942.2 | [link](https://huggingface.co/timm/maxvit_large_tf_512.in21k_ft_in1k)  |
+| maxvit_large_tf_384.in21k_ft_in1k  |   88   |         212   |  132.6 |   445.8 | [link](https://huggingface.co/timm/maxvit_large_tf_384.in21k_ft_in1k)  |
+| maxvit_base_tf_384.in21k_ft_in1k   |   87.9 |         119.6 |   73.8 |   332.9 | [link](https://huggingface.co/timm/maxvit_base_tf_384.in21k_ft_in1k)   |
+| maxvit_base_tf_512.in1k            |   86.6 |         119.9 |  138   |   704   | [link](https://huggingface.co/timm/maxvit_base_tf_512.in1k)            |
+| maxvit_large_tf_512.in1k           |   86.5 |         212.3 |  244.8 |   942.2 | [link](https://huggingface.co/timm/maxvit_large_tf_512.in1k)           |
+| maxvit_base_tf_384.in1k            |   86.3 |         119.6 |   73.8 |   332.9 | [link](https://huggingface.co/timm/maxvit_base_tf_384.in1k)            |
+| maxvit_large_tf_384.in1k           |   86.2 |         212   |  132.6 |   445.8 | [link](https://huggingface.co/timm/maxvit_large_tf_384.in1k)           |
+| maxvit_small_tf_512.in1k           |   86.1 |          69.1 |   67.3 |   383.8 | [link](https://huggingface.co/timm/maxvit_small_tf_512.in1k)           |
+| maxvit_tiny_tf_512.in1k            |   85.7 |          31   |   33.5 |   257.6 | [link](https://huggingface.co/timm/maxvit_tiny_tf_512.in1k)            |
+| maxvit_small_tf_384.in1k           |   85.5 |          69   |   35.9 |   183.6 | [link](https://huggingface.co/timm/maxvit_small_tf_384.in1k)           |
+| maxvit_tiny_tf_384.in1k            |   85.1 |          31   |   17.5 |   123.4 | [link](https://huggingface.co/timm/maxvit_tiny_tf_384.in1k)            |
+| maxvit_large_tf_224.in1k           |   84.9 |         211.8 |   43.7 |   127.4 | [link](https://huggingface.co/timm/maxvit_large_tf_224.in1k)           |
+| maxvit_base_tf_224.in1k            |   84.9 |         119.5 |   24   |    95   | [link](https://huggingface.co/timm/maxvit_base_tf_224.in1k)            |
+| maxvit_small_tf_224.in1k           |   84.4 |          68.9 |   11.7 |    53.2 | [link](https://huggingface.co/timm/maxvit_small_tf_224.in1k)           |
+| maxvit_tiny_tf_224.in1k            |   83.4 |          30.9 |    5.6 |    35.8 | [link](https://huggingface.co/timm/maxvit_tiny_tf_224.in1k)            |
+
 ### Oct 15, 2022
 * Train and validation script enhancements
 * Non-GPU (ie CPU) device support

benchmark.py

@@ -56,6 +56,7 @@ try:
 except ImportError as e:
     has_functorch = False
 
+has_compile = hasattr(torch, 'compile')
 
 if torch.cuda.is_available():
     torch.backends.cuda.matmul.allow_tf32 = True
@@ -74,8 +75,10 @@ parser.add_argument('--detail', action='store_true', default=False,
                     help='Provide train fwd/bwd/opt breakdown detail if True. Defaults to False')
 parser.add_argument('--no-retry', action='store_true', default=False,
                     help='Do not decay batch size and retry on error.')
-parser.add_argument('--results-file', default='', type=str, metavar='FILENAME',
+parser.add_argument('--results-file', default='', type=str,
                     help='Output csv file for validation results (summary)')
+parser.add_argument('--results-format', default='csv', type=str,
+                    help='Format for results file one of (csv, json) (default: csv).')
 parser.add_argument('--num-warm-iter', default=10, type=int,
                     metavar='N', help='Number of warmup iterations (default: 10)')
 parser.add_argument('--num-bench-iter', default=40, type=int,
@@ -106,13 +109,18 @@ parser.add_argument('--precision', default='float32', type=str,
                     help='Numeric precision. One of (amp, float32, float16, bfloat16, tf32)')
 parser.add_argument('--fuser', default='', type=str,
                     help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
+parser.add_argument('--fast-norm', default=False, action='store_true',
+                    help='enable experimental fast-norm')
+
+# codegen (model compilation) options
 scripting_group = parser.add_mutually_exclusive_group()
 scripting_group.add_argument('--torchscript', dest='torchscript', action='store_true',
-                    help='convert model torchscript for inference')
+                             help='convert model torchscript for inference')
+scripting_group.add_argument('--torchcompile', nargs='?', type=str, default=None, const='inductor',
+                             help="Enable compilation w/ specified backend (default: inductor).")
 scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
-                    help="Enable AOT Autograd support. (It's recommended to use this option with `--fuser nvfuser` together)")
-scripting_group.add_argument('--fast-norm', default=False, action='store_true',
-                    help='enable experimental fast-norm')
+                             help="Enable AOT Autograd optimization.")
+
 
 # train optimizer parameters
 parser.add_argument('--opt', default='sgd', type=str, metavar='OPTIMIZER',
@@ -205,6 +213,7 @@ class BenchmarkRunner:
             detail=False,
             device='cuda',
             torchscript=False,
+            torchcompile=None,
             aot_autograd=False,
             precision='float32',
             fuser='',
@@ -241,16 +250,22 @@ class BenchmarkRunner:
         _logger.info('Model %s created, param count: %d' % (model_name, self.param_count))
 
         data_config = resolve_data_config(kwargs, model=self.model, use_test_size=not use_train_size)
-        self.scripted = False
-        if torchscript:
-            self.model = torch.jit.script(self.model)
-            self.scripted = True
         self.input_size = data_config['input_size']
         self.batch_size = kwargs.pop('batch_size', 256)
 
-        if aot_autograd:
+        self.compiled = False
+        if torchscript:
+            self.model = torch.jit.script(self.model)
+            self.compiled = True
+        elif torchcompile:
+            assert has_compile, 'A version of torch w/ torch.compile() is required, possibly a nightly.'
+            torch._dynamo.reset()
+            self.model = torch.compile(self.model, backend=torchcompile)
+            self.compiled = True
+        elif aot_autograd:
             assert has_functorch, "functorch is needed for --aot-autograd"
             self.model = memory_efficient_fusion(self.model)
+            self.compiled = True
 
         self.example_inputs = None
         self.num_warm_iter = num_warm_iter
@@ -322,7 +337,7 @@ class InferenceBenchmarkRunner(BenchmarkRunner):
             param_count=round(self.param_count / 1e6, 2),
         )
 
-        retries = 0 if self.scripted else 2  # skip profiling if model is scripted
+        retries = 0 if self.compiled else 2  # skip profiling if model is scripted
         while retries:
             retries -= 1
             try:
@@ -620,7 +635,6 @@ def main():
         model_cfgs = [(n, None) for n in model_names]
 
     if len(model_cfgs):
-        results_file = args.results_file or './benchmark.csv'
         _logger.info('Running bulk validation on these pretrained models: {}'.format(', '.join(model_names)))
         results = []
         try:
@@ -641,22 +655,30 @@ def main():
             sort_key = 'infer_gmacs'
         results = filter(lambda x: sort_key in x, results)
         results = sorted(results, key=lambda x: x[sort_key], reverse=True)
-        if len(results):
-            write_results(results_file, results)
     else:
         results = benchmark(args)
 
+    if args.results_file:
+        write_results(args.results_file, results, format=args.results_format)
+
     # output results in JSON to stdout w/ delimiter for runner script
     print(f'--result\n{json.dumps(results, indent=4)}')
 
 
-def write_results(results_file, results):
+def write_results(results_file, results, format='csv'):
     with open(results_file, mode='w') as cf:
-        dw = csv.DictWriter(cf, fieldnames=results[0].keys())
-        dw.writeheader()
-        for r in results:
-            dw.writerow(r)
-        cf.flush()
+        if format == 'json':
+            json.dump(results, cf, indent=4)
+        else:
+            if not isinstance(results, (list, tuple)):
+                results = [results]
+            if not results:
+                return
+            dw = csv.DictWriter(cf, fieldnames=results[0].keys())
+            dw.writeheader()
+            for r in results:
+                dw.writerow(r)
+            cf.flush()
 
 
 if __name__ == '__main__':

inference.py

@@ -8,41 +8,89 @@ Hacked together by / Copyright 2020 Ross Wightman (https://github.com/rwightman)
 import os
 import time
 import argparse
+import json
 import logging
+from contextlib import suppress
+from functools import partial
+
 import numpy as np
+import pandas as pd
 import torch
 
-from timm.models import create_model, apply_test_time_pool
-from timm.data import ImageDataset, create_loader, resolve_data_config
-from timm.utils import AverageMeter, setup_default_logging
+from timm.models import create_model, apply_test_time_pool, load_checkpoint
+from timm.data import create_dataset, create_loader, resolve_data_config
+from timm.utils import AverageMeter, setup_default_logging, set_jit_fuser
+
+
+
+try:
+    from apex import amp
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+has_native_amp = False
+try:
+    if getattr(torch.cuda.amp, 'autocast') is not None:
+        has_native_amp = True
+except AttributeError:
+    pass
+
+try:
+    from functorch.compile import memory_efficient_fusion
+    has_functorch = True
+except ImportError as e:
+    has_functorch = False
+
+has_compile = hasattr(torch, 'compile')
+
+
+_FMT_EXT = {
+    'json': '.json',
+    'json-split': '.json',
+    'parquet': '.parquet',
+    'csv': '.csv',
+}
 
 torch.backends.cudnn.benchmark = True
 _logger = logging.getLogger('inference')
 
 
 parser = argparse.ArgumentParser(description='PyTorch ImageNet Inference')
-parser.add_argument('data', metavar='DIR',
-                    help='path to dataset')
-parser.add_argument('--output_dir', metavar='DIR', default='./',
-                    help='path to output files')
-parser.add_argument('--model', '-m', metavar='MODEL', default='dpn92',
-                    help='model architecture (default: dpn92)')
+parser.add_argument('data', nargs='?', metavar='DIR', const=None,
+                    help='path to dataset (*deprecated*, use --data-dir)')
+parser.add_argument('--data-dir', metavar='DIR',
+                    help='path to dataset (root dir)')
+parser.add_argument('--dataset', metavar='NAME', default='',
+                    help='dataset type + name ("<type>/<name>") (default: ImageFolder or ImageTar if empty)')
+parser.add_argument('--split', metavar='NAME', default='validation',
+                    help='dataset split (default: validation)')
+parser.add_argument('--model', '-m', metavar='MODEL', default='resnet50',
+                    help='model architecture (default: resnet50)')
 parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
                     help='number of data loading workers (default: 2)')
 parser.add_argument('-b', '--batch-size', default=256, type=int,
                     metavar='N', help='mini-batch size (default: 256)')
 parser.add_argument('--img-size', default=None, type=int,
-                    metavar='N', help='Input image dimension')
+                    metavar='N', help='Input image dimension, uses model default if empty')
 parser.add_argument('--input-size', default=None, nargs=3, type=int,
                     metavar='N N N', help='Input all image dimensions (d h w, e.g. --input-size 3 224 224), uses model default if empty')
+parser.add_argument('--use-train-size', action='store_true', default=False,
+                    help='force use of train input size, even when test size is specified in pretrained cfg')
+parser.add_argument('--crop-pct', default=None, type=float,
+                    metavar='N', help='Input image center crop pct')
+parser.add_argument('--crop-mode', default=None, type=str,
+                    metavar='N', help='Input image crop mode (squash, border, center). Model default if None.')
 parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
                     help='Override mean pixel value of dataset')
-parser.add_argument('--std', type=float, nargs='+', default=None, metavar='STD',
+parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
                     help='Override std deviation of of dataset')
 parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
                     help='Image resize interpolation type (overrides model)')
-parser.add_argument('--num-classes', type=int, default=1000,
+parser.add_argument('--num-classes', type=int, default=None,
                     help='Number classes in dataset')
+parser.add_argument('--class-map', default='', type=str, metavar='FILENAME',
+                    help='path to class to idx mapping file (default: "")')
 parser.add_argument('--log-freq', default=10, type=int,
                     metavar='N', help='batch logging frequency (default: 10)')
 parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
@@ -51,10 +99,49 @@ parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                     help='use pre-trained model')
 parser.add_argument('--num-gpu', type=int, default=1,
                     help='Number of GPUS to use')
-parser.add_argument('--no-test-pool', dest='no_test_pool', action='store_true',
-                    help='disable test time pool')
-parser.add_argument('--topk', default=5, type=int,
+parser.add_argument('--test-pool', dest='test_pool', action='store_true',
+                    help='enable test time pool')
+parser.add_argument('--channels-last', action='store_true', default=False,
+                    help='Use channels_last memory layout')
+parser.add_argument('--device', default='cuda', type=str,
+                    help="Device (accelerator) to use.")
+parser.add_argument('--amp', action='store_true', default=False,
+                    help='use Native AMP for mixed precision training')
+parser.add_argument('--amp-dtype', default='float16', type=str,
+                    help='lower precision AMP dtype (default: float16)')
+parser.add_argument('--fuser', default='', type=str,
+                    help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
+
+scripting_group = parser.add_mutually_exclusive_group()
+scripting_group.add_argument('--torchscript', default=False, action='store_true',
+                             help='torch.jit.script the full model')
+scripting_group.add_argument('--torchcompile', nargs='?', type=str, default=None, const='inductor',
+                             help="Enable compilation w/ specified backend (default: inductor).")
+scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
+                             help="Enable AOT Autograd support.")
+
+parser.add_argument('--results-dir',type=str, default=None,
+                    help='folder for output results')
+parser.add_argument('--results-file', type=str, default=None,
+                    help='results filename (relative to results-dir)')
+parser.add_argument('--results-format', type=str, nargs='+', default=['csv'],
+                    help='results format (one of "csv", "json", "json-split", "parquet")')
+parser.add_argument('--results-separate-col', action='store_true', default=False,
+                    help='separate output columns per result index.')
+parser.add_argument('--topk', default=1, type=int,
                     metavar='N', help='Top-k to output to CSV')
+parser.add_argument('--fullname', action='store_true', default=False,
+                    help='use full sample name in output (not just basename).')
+parser.add_argument('--filename-col', default='filename',
+                    help='name for filename / sample name column')
+parser.add_argument('--index-col', default='index',
+                    help='name for output indices column(s)')
+parser.add_argument('--output-col', default=None,
+                    help='name for logit/probs output column(s)')
+parser.add_argument('--output-type', default='prob',
+                    help='output type colum ("prob" for probabilities, "logit" for raw logits)')
+parser.add_argument('--exclude-output', action='store_true', default=False,
+                    help='exclude logits/probs from results, just indices. topk must be set !=0.')
 
 
 def main():
@@ -63,48 +150,104 @@ def main():
     # might as well try to do something useful...
     args.pretrained = args.pretrained or not args.checkpoint
 
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.benchmark = True
+
+    device = torch.device(args.device)
+
+    # resolve AMP arguments based on PyTorch / Apex availability
+    amp_autocast = suppress
+    if args.amp:
+        assert has_native_amp, 'Please update PyTorch to a version with native AMP (or use APEX).'
+        assert args.amp_dtype in ('float16', 'bfloat16')
+        amp_dtype = torch.bfloat16 if args.amp_dtype == 'bfloat16' else torch.float16
+        amp_autocast = partial(torch.autocast, device_type=device.type, dtype=amp_dtype)
+        _logger.info('Running inference in mixed precision with native PyTorch AMP.')
+    else:
+        _logger.info('Running inference in float32. AMP not enabled.')
+
+    if args.fuser:
+        set_jit_fuser(args.fuser)
+
     # create model
     model = create_model(
         args.model,
         num_classes=args.num_classes,
         in_chans=3,
         pretrained=args.pretrained,
-        checkpoint_path=args.checkpoint)
+        checkpoint_path=args.checkpoint,
+    )
+    if args.num_classes is None:
+        assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'
+        args.num_classes = model.num_classes
+
+    _logger.info(
+        f'Model {args.model} created, param count: {sum([m.numel() for m in model.parameters()])}')
 
-    _logger.info('Model %s created, param count: %d' %
-                 (args.model, sum([m.numel() for m in model.parameters()])))
+    data_config = resolve_data_config(vars(args), model=model)
+    test_time_pool = False
+    if args.test_pool:
+        model, test_time_pool = apply_test_time_pool(model, data_config)
+
+    model = model.to(device)
+    model.eval()
+    if args.channels_last:
+        model = model.to(memory_format=torch.channels_last)
 
-    config = resolve_data_config(vars(args), model=model)
-    model, test_time_pool = (model, False) if args.no_test_pool else apply_test_time_pool(model, config)
+    if args.torchscript:
+        model = torch.jit.script(model)
+    elif args.torchcompile:
+        assert has_compile, 'A version of torch w/ torch.compile() is required for --compile, possibly a nightly.'
+        torch._dynamo.reset()
+        model = torch.compile(model, backend=args.torchcompile)
+    elif args.aot_autograd:
+        assert has_functorch, "functorch is needed for --aot-autograd"
+        model = memory_efficient_fusion(model)
 
     if args.num_gpu > 1:
-        model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda()
-    else:
-        model = model.cuda()
+        model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))
+
+    root_dir = args.data or args.data_dir
+    dataset = create_dataset(
+        root=root_dir,
+        name=args.dataset,
+        split=args.split,
+        class_map=args.class_map,
+    )
+
+    if test_time_pool:
+        data_config['crop_pct'] = 1.0
 
+    workers = 1 if 'tfds' in args.dataset or 'wds' in args.dataset else args.workers
     loader = create_loader(
-        ImageDataset(args.data),
-        input_size=config['input_size'],
+        dataset,
         batch_size=args.batch_size,
         use_prefetcher=True,
-        interpolation=config['interpolation'],
-        mean=config['mean'],
-        std=config['std'],
-        num_workers=args.workers,
-        crop_pct=1.0 if test_time_pool else config['crop_pct'])
-
-    model.eval()
+        num_workers=workers,
+        **data_config,
+    )
 
-    k = min(args.topk, args.num_classes)
+    top_k = min(args.topk, args.num_classes)
     batch_time = AverageMeter()
     end = time.time()
-    topk_ids = []
+    all_indices = []
+    all_outputs = []
+    use_probs = args.output_type == 'prob'
     with torch.no_grad():
         for batch_idx, (input, _) in enumerate(loader):
-            input = input.cuda()
-            labels = model(input)
-            topk = labels.topk(k)[1]
-            topk_ids.append(topk.cpu().numpy())
+
+            with amp_autocast():
+                output = model(input)
+
+            if use_probs:
+                output = output.softmax(-1)
+
+            if top_k:
+                output, indices = output.topk(top_k)
+                all_indices.append(indices.cpu().numpy())
+
+            all_outputs.append(output.cpu().numpy())
 
             # measure elapsed time
             batch_time.update(time.time() - end)
@@ -114,13 +257,61 @@ def main():
                 _logger.info('Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
                     batch_idx, len(loader), batch_time=batch_time))
 
-    topk_ids = np.concatenate(topk_ids, axis=0)
+    all_indices = np.concatenate(all_indices, axis=0) if all_indices else None
+    all_outputs = np.concatenate(all_outputs, axis=0).astype(np.float32)
+    filenames = loader.dataset.filenames(basename=not args.fullname)
 
-    with open(os.path.join(args.output_dir, './topk_ids.csv'), 'w') as out_file:
-        filenames = loader.dataset.filenames(basename=True)
-        for filename, label in zip(filenames, topk_ids):
-            out_file.write('{0},{1}\n'.format(
-                filename, ','.join([ str(v) for v in label])))
+    output_col = args.output_col or ('prob' if use_probs else 'logit')
+    data_dict = {args.filename_col: filenames}
+    if args.results_separate_col and all_outputs.shape[-1] > 1:
+        if all_indices is not None:
+            for i in range(all_indices.shape[-1]):
+                data_dict[f'{args.index_col}_{i}'] = all_indices[:, i]
+        for i in range(all_outputs.shape[-1]):
+            data_dict[f'{output_col}_{i}'] = all_outputs[:, i]
+    else:
+        if all_indices is not None:
+            if all_indices.shape[-1] == 1:
+                all_indices = all_indices.squeeze(-1)
+            data_dict[args.index_col] = list(all_indices)
+        if all_outputs.shape[-1] == 1:
+            all_outputs = all_outputs.squeeze(-1)
+        data_dict[output_col] = list(all_outputs)
+
+    df = pd.DataFrame(data=data_dict)
+
+    results_filename = args.results_file
+    if results_filename:
+        filename_no_ext, ext = os.path.splitext(results_filename)[-1]
+        if ext and ext in _FMT_EXT.values():
+            # if filename provided with one of expected ext,
+            # remove it as it will be added back
+            results_filename = filename_no_ext
+    else:
+        # base default filename on model name + img-size
+        img_size = data_config["input_size"][1]
+        results_filename = f'{args.model}-{img_size}'
+
+    if args.results_dir:
+        results_filename = os.path.join(args.results_dir, results_filename)
+
+    for fmt in args.results_format:
+        save_results(df, results_filename, fmt)
+
+    print(f'--result')
+    print(json.dumps(dict(filename=results_filename)))
+
+
+def save_results(df, results_filename, results_format='csv', filename_col='filename'):
+    results_filename += _FMT_EXT[results_format]
+    if results_format == 'parquet':
+        df.set_index(filename_col).to_parquet(results_filename)
+    elif results_format == 'json':
+        df.to_json(results_filename, lines=True, orient='records')
+    elif results_format == 'json-split':
+        df.to_json(results_filename, indent=4, orient='split', index=False)
+    else:
+        df.to_csv(results_filename, index=False)
 
 
 if __name__ == '__main__':

tests/test_models.py

@@ -13,8 +13,7 @@ except ImportError:
     has_fx_feature_extraction = False
 
 import timm
-from timm import list_models, create_model, set_scriptable, has_pretrained_cfg_key, is_pretrained_cfg_key, \
-    get_pretrained_cfg_value
+from timm import list_models, create_model, set_scriptable, get_pretrained_cfg_value
 from timm.models.fx_features import _leaf_modules, _autowrap_functions    
 
 if hasattr(torch._C, '_jit_set_profiling_executor'):

timm/__init__.py

@@ -1,4 +1,4 @@
 from .version import __version__
-from .models import create_model, list_models, is_model, list_modules, model_entrypoint, \
-    is_scriptable, is_exportable, set_scriptable, set_exportable, has_pretrained_cfg_key, is_pretrained_cfg_key, \
-    get_pretrained_cfg_value, is_model_pretrained
+from .models import create_model, list_models, list_pretrained, is_model, list_modules, model_entrypoint, \
+    is_scriptable, is_exportable, set_scriptable, set_exportable, \
+    is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value

timm/data/config.py

@@ -5,75 +5,89 @@ from .constants import *
 _logger = logging.getLogger(__name__)
 
 
-def resolve_data_config(args, default_cfg={}, model=None, use_test_size=False, verbose=False):
+def resolve_data_config(
+        args,
+        default_cfg=None,
+        model=None,
+        use_test_size=False,
+        verbose=False
+):
     new_config = {}
-    default_cfg = default_cfg
+    default_cfg = default_cfg or {}
     if not default_cfg and model is not None and hasattr(model, 'default_cfg'):
         default_cfg = model.default_cfg
 
     # Resolve input/image size
     in_chans = 3
-    if 'chans' in args and args['chans'] is not None:
+    if args.get('chans', None) is not None:
         in_chans = args['chans']
 
     input_size = (in_chans, 224, 224)
-    if 'input_size' in args and args['input_size'] is not None:
+    if args.get('input_size', None) is not None:
         assert isinstance(args['input_size'], (tuple, list))
         assert len(args['input_size']) == 3
         input_size = tuple(args['input_size'])
         in_chans = input_size[0]  # input_size overrides in_chans
-    elif 'img_size' in args and args['img_size'] is not None:
+    elif args.get('img_size', None) is not None:
         assert isinstance(args['img_size'], int)
         input_size = (in_chans, args['img_size'], args['img_size'])
     else:
-        if use_test_size and 'test_input_size' in default_cfg:
+        if use_test_size and default_cfg.get('test_input_size', None) is not None:
             input_size = default_cfg['test_input_size']
-        elif 'input_size' in default_cfg:
+        elif default_cfg.get('input_size', None) is not None:
             input_size = default_cfg['input_size']
     new_config['input_size'] = input_size
 
     # resolve interpolation method
     new_config['interpolation'] = 'bicubic'
-    if 'interpolation' in args and args['interpolation']:
+    if args.get('interpolation', None):
         new_config['interpolation'] = args['interpolation']
-    elif 'interpolation' in default_cfg:
+    elif default_cfg.get('interpolation', None):
         new_config['interpolation'] = default_cfg['interpolation']
 
     # resolve dataset + model mean for normalization
     new_config['mean'] = IMAGENET_DEFAULT_MEAN
-    if 'mean' in args and args['mean'] is not None:
+    if args.get('mean', None) is not None:
         mean = tuple(args['mean'])
         if len(mean) == 1:
             mean = tuple(list(mean) * in_chans)
         else:
             assert len(mean) == in_chans
         new_config['mean'] = mean
-    elif 'mean' in default_cfg:
+    elif default_cfg.get('mean', None):
         new_config['mean'] = default_cfg['mean']
 
     # resolve dataset + model std deviation for normalization
     new_config['std'] = IMAGENET_DEFAULT_STD
-    if 'std' in args and args['std'] is not None:
+    if args.get('std', None) is not None:
         std = tuple(args['std'])
         if len(std) == 1:
             std = tuple(list(std) * in_chans)
         else:
             assert len(std) == in_chans
         new_config['std'] = std
-    elif 'std' in default_cfg:
+    elif default_cfg.get('std', None):
         new_config['std'] = default_cfg['std']
 
-    # resolve default crop percentage
+    # resolve default inference crop
     crop_pct = DEFAULT_CROP_PCT
-    if 'crop_pct' in args and args['crop_pct'] is not None:
+    if args.get('crop_pct', None):
         crop_pct = args['crop_pct']
     else:
-        if use_test_size and 'test_crop_pct' in default_cfg:
+        if use_test_size and default_cfg.get('test_crop_pct', None):
             crop_pct = default_cfg['test_crop_pct']
-        elif 'crop_pct' in default_cfg:
+        elif default_cfg.get('crop_pct', None):
             crop_pct = default_cfg['crop_pct']
     new_config['crop_pct'] = crop_pct
 
+    # resolve default crop percentage
+    crop_mode = DEFAULT_CROP_MODE
+    if args.get('crop_mode', None):
+        crop_mode = args['crop_mode']
+    elif default_cfg.get('crop_mode', None):
+        crop_mode = default_cfg['crop_mode']
+    new_config['crop_mode'] = crop_mode
+
     if verbose:
         _logger.info('Data processing configuration for current model + dataset:')
         for n, v in new_config.items():

timm/data/constants.py

@@ -1,4 +1,5 @@
 DEFAULT_CROP_PCT = 0.875
+DEFAULT_CROP_MODE = 'center'
 IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
 IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5)

timm/data/loader.py

@@ -211,6 +211,7 @@ def create_loader(
         num_workers=1,
         distributed=False,
         crop_pct=None,
+        crop_mode=None,
         collate_fn=None,
         pin_memory=False,
         fp16=False,  # deprecated, use img_dtype
@@ -240,6 +241,7 @@ def create_loader(
         mean=mean,
         std=std,
         crop_pct=crop_pct,
+        crop_mode=crop_mode,
         tf_preprocessing=tf_preprocessing,
         re_prob=re_prob,
         re_mode=re_mode,

timm/data/readers/reader_tfds.py

@@ -43,6 +43,15 @@ SHUFFLE_SIZE = int(os.environ.get('TFDS_SHUFFLE_SIZE', 8192))  # samples to shuf
 PREFETCH_SIZE = int(os.environ.get('TFDS_PREFETCH_SIZE', 2048))  # samples to prefetch
 
 
+@tfds.decode.make_decoder()
+def decode_example(serialized_image, feature, dct_method='INTEGER_ACCURATE'):
+    return tf.image.decode_jpeg(
+        serialized_image,
+        channels=3,
+        dct_method=dct_method,
+    )
+
+
 def even_split_indices(split, n, num_samples):
     partitions = [round(i * num_samples / n) for i in range(n + 1)]
     return [f"{split}[{partitions[i]}:{partitions[i + 1]}]" for i in range(n)]
@@ -242,6 +251,7 @@ class ReaderTfds(Reader):
         ds = self.builder.as_dataset(
             split=self.subsplit or self.split,
             shuffle_files=self.is_training,
+            decoders=dict(image=decode_example()),
             read_config=read_config,
         )
         # avoid overloading threading w/ combo of TF ds threads + PyTorch workers

timm/data/tf_preprocessing.py

@@ -22,12 +22,13 @@ Hacked together by / Copyright 2020 Ross Wightman
 # limitations under the License.
 # ==============================================================================
 """ImageNet preprocessing for MnasNet."""
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 
 IMAGE_SIZE = 224
 CROP_PADDING = 32
 
+tf.compat.v1.disable_eager_execution()
 
 def distorted_bounding_box_crop(image_bytes,
                                 bbox,

timm/data/transforms.py

@@ -1,3 +1,9 @@
+import math
+import numbers
+import random
+import warnings
+from typing import List, Sequence
+
 import torch
 import torchvision.transforms.functional as F
 try:
@@ -6,9 +12,6 @@ try:
 except ImportError:
     has_interpolation_mode = False
 from PIL import Image
-import warnings
-import math
-import random
 import numpy as np
 
 
@@ -96,6 +99,19 @@ def interp_mode_to_str(mode):
 _RANDOM_INTERPOLATION = (str_to_interp_mode('bilinear'), str_to_interp_mode('bicubic'))
 
 
+def _setup_size(size, error_msg):
+    if isinstance(size, numbers.Number):
+        return int(size), int(size)
+
+    if isinstance(size, Sequence) and len(size) == 1:
+        return size[0], size[0]
+
+    if len(size) != 2:
+        raise ValueError(error_msg)
+
+    return size
+
+
 class RandomResizedCropAndInterpolation:
     """Crop the given PIL Image to random size and aspect ratio with random interpolation.
 
@@ -195,3 +211,132 @@ class RandomResizedCropAndInterpolation:
         format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
         format_string += ', interpolation={0})'.format(interpolate_str)
         return format_string
+
+
+def center_crop_or_pad(img: torch.Tensor, output_size: List[int], fill=0) -> torch.Tensor:
+    """Center crops and/or pads the given image.
+    If the image is torch Tensor, it is expected
+    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
+
+    Args:
+        img (PIL Image or Tensor): Image to be cropped.
+        output_size (sequence or int): (height, width) of the crop box. If int or sequence with single int,
+            it is used for both directions.
+        fill (int, Tuple[int]): Padding color
+
+    Returns:
+        PIL Image or Tensor: Cropped image.
+    """
+    if isinstance(output_size, numbers.Number):
+        output_size = (int(output_size), int(output_size))
+    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
+        output_size = (output_size[0], output_size[0])
+
+    _, image_height, image_width = F.get_dimensions(img)
+    crop_height, crop_width = output_size
+
+    if crop_width > image_width or crop_height > image_height:
+        padding_ltrb = [
+            (crop_width - image_width) // 2 if crop_width > image_width else 0,
+            (crop_height - image_height) // 2 if crop_height > image_height else 0,
+            (crop_width - image_width + 1) // 2 if crop_width > image_width else 0,
+            (crop_height - image_height + 1) // 2 if crop_height > image_height else 0,
+        ]
+        img = F.pad(img, padding_ltrb, fill=fill)
+        _, image_height, image_width = F.get_dimensions(img)
+        if crop_width == image_width and crop_height == image_height:
+            return img
+
+    crop_top = int(round((image_height - crop_height) / 2.0))
+    crop_left = int(round((image_width - crop_width) / 2.0))
+    return F.crop(img, crop_top, crop_left, crop_height, crop_width)
+
+
+class CenterCropOrPad(torch.nn.Module):
+    """Crops the given image at the center.
+    If the image is torch Tensor, it is expected
+    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
+
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
+    """
+
+    def __init__(self, size, fill=0):
+        super().__init__()
+        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
+        self.fill = fill
+
+    def forward(self, img):
+        """
+        Args:
+            img (PIL Image or Tensor): Image to be cropped.
+
+        Returns:
+            PIL Image or Tensor: Cropped image.
+        """
+        return center_crop_or_pad(img, self.size, fill=self.fill)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+class ResizeKeepRatio:
+    """ Resize and Keep Ratio
+    """
+
+    def __init__(
+            self,
+            size,
+            longest=0.,
+            interpolation='bilinear',
+            fill=0,
+    ):
+        if isinstance(size, (list, tuple)):
+            self.size = tuple(size)
+        else:
+            self.size = (size, size)
+        self.interpolation = str_to_interp_mode(interpolation)
+        self.longest = float(longest)
+        self.fill = fill
+
+    @staticmethod
+    def get_params(img, target_size, longest):
+        """Get parameters
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            target_size (Tuple[int, int]): Size of output
+        Returns:
+            tuple: params (h, w) and (l, r, t, b) to be passed to ``resize`` and ``pad`` respectively
+        """
+        source_size = img.size[::-1]  # h, w
+        h, w = source_size
+        target_h, target_w = target_size
+        ratio_h = h / target_h
+        ratio_w = w / target_w
+        ratio = max(ratio_h, ratio_w) * longest + min(ratio_h, ratio_w) * (1. - longest)
+        size = [round(x / ratio) for x in source_size]
+        return size
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be cropped and resized.
+
+        Returns:
+            PIL Image: Resized, padded to at least target size, possibly cropped to exactly target size
+        """
+        size = self.get_params(img, self.size, self.longest)
+        img = F.resize(img, size, self.interpolation)
+        return img
+
+    def __repr__(self):
+        interpolate_str = interp_mode_to_str(self.interpolation)
+        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
+        format_string += f', interpolation={interpolate_str})'
+        format_string += f', longest={self.longest:.3f})'
+        return format_string

timm/data/transforms_factory.py

@@ -10,7 +10,8 @@ from torchvision import transforms
 
 from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, DEFAULT_CROP_PCT
 from timm.data.auto_augment import rand_augment_transform, augment_and_mix_transform, auto_augment_transform
-from timm.data.transforms import str_to_interp_mode, str_to_pil_interp, RandomResizedCropAndInterpolation, ToNumpy
+from timm.data.transforms import str_to_interp_mode, str_to_pil_interp, RandomResizedCropAndInterpolation,\
+    ResizeKeepRatio, CenterCropOrPad, ToNumpy
 from timm.data.random_erasing import RandomErasing
 
 
@@ -130,26 +131,49 @@ def transforms_imagenet_train(
 def transforms_imagenet_eval(
         img_size=224,
         crop_pct=None,
+        crop_mode=None,
         interpolation='bilinear',
         use_prefetcher=False,
         mean=IMAGENET_DEFAULT_MEAN,
-        std=IMAGENET_DEFAULT_STD):
+        std=IMAGENET_DEFAULT_STD
+):
     crop_pct = crop_pct or DEFAULT_CROP_PCT
 
     if isinstance(img_size, (tuple, list)):
         assert len(img_size) == 2
-        if img_size[-1] == img_size[-2]:
-            # fall-back to older behaviour so Resize scales to shortest edge if target is square
-            scale_size = int(math.floor(img_size[0] / crop_pct))
-        else:
-            scale_size = tuple([int(x / crop_pct) for x in img_size])
+        scale_size = tuple([math.floor(x / crop_pct) for x in img_size])
     else:
-        scale_size = int(math.floor(img_size / crop_pct))
+        scale_size = math.floor(img_size / crop_pct)
+        scale_size = (scale_size, scale_size)
+
+    if crop_mode == 'squash':
+        # squash mode scales each edge to 1/pct of target, then crops
+        # aspect ratio is not preserved, no img lost if crop_pct == 1.0
+        tfl = [
+            transforms.Resize(scale_size, interpolation=str_to_interp_mode(interpolation)),
+            transforms.CenterCrop(img_size),
+        ]
+    elif crop_mode == 'border':
+        # scale the longest edge of image to 1/pct of target edge, add borders to pad, then crop
+        # no image lost if crop_pct == 1.0
+        fill = [round(255 * v) for v in mean]
+        tfl = [
+            ResizeKeepRatio(scale_size, interpolation=interpolation, longest=1.0),
+            CenterCropOrPad(img_size, fill=fill),
+        ]
+    else:
+        # default crop model is center
+        # aspect ratio is preserved, crops center within image, no borders are added, image is lost
+        if scale_size[0] == scale_size[1]:
+            # simple case, use torchvision built-in Resize w/ shortest edge mode (scalar size arg)
+            tfl = [
+                transforms.Resize(scale_size[0], interpolation=str_to_interp_mode(interpolation))
+            ]
+        else:
+            # resize shortest edge to matching target dim for non-square target
+            tfl = [ResizeKeepRatio(scale_size)]
+        tfl += [transforms.CenterCrop(img_size)]
 
-    tfl = [
-        transforms.Resize(scale_size, interpolation=str_to_interp_mode(interpolation)),
-        transforms.CenterCrop(img_size),
-    ]
     if use_prefetcher:
         # prefetcher and collate will handle tensor conversion and norm
         tfl += [ToNumpy()]
@@ -157,8 +181,9 @@ def transforms_imagenet_eval(
         tfl += [
             transforms.ToTensor(),
             transforms.Normalize(
-                     mean=torch.tensor(mean),
-                     std=torch.tensor(std))
+                mean=torch.tensor(mean),
+                std=torch.tensor(std),
+            )
         ]
 
     return transforms.Compose(tfl)
@@ -183,6 +208,7 @@ def create_transform(
         re_count=1,
         re_num_splits=0,
         crop_pct=None,
+        crop_mode=None,
         tf_preprocessing=False,
         separate=False):
 
@@ -204,7 +230,8 @@ def create_transform(
                 interpolation=interpolation,
                 use_prefetcher=use_prefetcher,
                 mean=mean,
-                std=std)
+                std=std,
+            )
         elif is_training:
             transform = transforms_imagenet_train(
                 img_size,
@@ -222,7 +249,8 @@ def create_transform(
                 re_mode=re_mode,
                 re_count=re_count,
                 re_num_splits=re_num_splits,
-                separate=separate)
+                separate=separate,
+            )
         else:
             assert not separate, "Separate transforms not supported for validation preprocessing"
             transform = transforms_imagenet_eval(
@@ -231,6 +259,8 @@ def create_transform(
                 use_prefetcher=use_prefetcher,
                 mean=mean,
                 std=std,
-                crop_pct=crop_pct)
+                crop_pct=crop_pct,
+                crop_mode=crop_mode,
+            )
 
     return transform

timm/models/__init__.py

@@ -70,5 +70,6 @@ from .layers import TestTimePoolHead, apply_test_time_pool
 from .layers import convert_splitbn_model, convert_sync_batchnorm
 from .layers import is_scriptable, is_exportable, set_scriptable, set_exportable, is_no_jit, set_no_jit
 from .layers import set_fast_norm
-from .registry import register_model, model_entrypoint, list_models, is_model, list_modules, is_model_in_modules,\
-    is_model_pretrained, get_pretrained_cfg, has_pretrained_cfg_key, is_pretrained_cfg_key, get_pretrained_cfg_value
+from .pretrained import PretrainedCfg, filter_pretrained_cfg, generate_default_cfgs, split_model_name_tag
+from .registry import register_model, model_entrypoint, list_models, list_pretrained, is_model, list_modules,\
+    is_model_in_modules, is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value

timm/models/convnext.py

@@ -21,111 +21,13 @@ from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from .helpers import named_apply, build_model_with_cfg, checkpoint_seq
 from .layers import trunc_normal_, SelectAdaptivePool2d, DropPath, ConvMlp, Mlp, LayerNorm2d, LayerNorm, \
     create_conv2d, get_act_layer, make_divisible, to_ntuple
+from .pretrained import generate_default_cfgs
 from .registry import register_model
 
 
 __all__ = ['ConvNeXt']  # model_registry will add each entrypoint fn to this
 
 
-def _cfg(url='', **kwargs):
-    return {
-        'url': url,
-        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
-        'crop_pct': 0.875, 'interpolation': 'bicubic',
-        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
-        'first_conv': 'stem.0', 'classifier': 'head.fc',
-        **kwargs
-    }
-
-
-default_cfgs = dict(
-    # timm specific variants
-    convnext_atto=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_atto_d2-01bb0f51.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=0.95),
-    convnext_atto_ols=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_atto_ols_a2-78d1c8f3.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=0.95),
-    convnext_femto=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_femto_d1-d71d5b4c.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=0.95),
-    convnext_femto_ols=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_femto_ols_d1-246bf2ed.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=0.95),
-    convnext_pico=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_pico_d1-10ad7f0d.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=0.95),
-    convnext_pico_ols=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_pico_ols_d1-611f0ca7.pth',
-        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_nano=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_nano_d1h-7eb4bdea.pth',
-        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_nano_ols=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_nano_ols_d1h-ae424a9a.pth',
-        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_tiny_hnf=_cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_tiny_hnf_a2h-ab7e9df2.pth',
-        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
-
-    convnext_tiny=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_small=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_base=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_large=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-
-    convnext_tiny_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_1k_224.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_small_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_1k_224.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_base_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_large_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-    convnext_xlarge_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth',
-        test_input_size=(3, 288, 288), test_crop_pct=1.0),
-
-    convnext_tiny_384_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_1k_384.pth',
-        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0),
-    convnext_small_384_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_1k_384.pth',
-        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0),
-    convnext_base_384_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth',
-        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0),
-    convnext_large_384_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth',
-        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0),
-    convnext_xlarge_384_in22ft1k=_cfg(
-        url='https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth',
-        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0),
-
-    convnext_tiny_in22k=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth", num_classes=21841),
-    convnext_small_in22k=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth", num_classes=21841),
-    convnext_base_in22k=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth", num_classes=21841),
-    convnext_large_in22k=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth", num_classes=21841),
-    convnext_xlarge_in22k=_cfg(
-        url="https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth", num_classes=21841),
-)
-
-
 class ConvNeXtBlock(nn.Module):
     """ ConvNeXt Block
     There are two equivalent implementations:
@@ -459,6 +361,107 @@ def _create_convnext(variant, pretrained=False, **kwargs):
     return model
 
 
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
+        'crop_pct': 0.875, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
+        'first_conv': 'stem.0', 'classifier': 'head.fc',
+        **kwargs
+    }
+
+
+default_cfgs = generate_default_cfgs({
+    # timm specific variants
+    'convnext_atto.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_atto_d2-01bb0f51.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=0.95),
+    'convnext_atto_ols.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_atto_ols_a2-78d1c8f3.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=0.95),
+    'convnext_femto.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_femto_d1-d71d5b4c.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=0.95),
+    'convnext_femto_ols.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_femto_ols_d1-246bf2ed.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=0.95),
+    'convnext_pico.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_pico_d1-10ad7f0d.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=0.95),
+    'convnext_pico_ols.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_pico_ols_d1-611f0ca7.pth',
+        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_nano.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_nano_d1h-7eb4bdea.pth',
+        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_nano_ols.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_nano_ols_d1h-ae424a9a.pth',
+        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_tiny_hnf.timm_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_tiny_hnf_a2h-ab7e9df2.pth',
+        crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
+
+    'convnext_tiny.fb_in1k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_small.fb_in1k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_base.fb_in1k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_large.fb_in1k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_xlarge.untrained': _cfg(),
+
+    'convnext_tiny.fb_in22k_ft_in1k': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_1k_224.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_small.fb_in22k_ft_in1k': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_1k_224.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_base.fb_in22k_ft_in1k': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_large.fb_in22k_ft_in1k': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+    'convnext_xlarge.fb_in22k_ft_in1k': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth',
+        test_input_size=(3, 288, 288), test_crop_pct=1.0),
+
+    'convnext_tiny.fb_in22k_ft_in1k_384': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_1k_384.pth',
+        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
+    'convnext_small..fb_in22k_ft_in1k_384': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_1k_384.pth',
+        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
+    'convnext_base.fb_in22k_ft_in1k_384': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth',
+        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
+    'convnext_large.fb_in22k_ft_in1k_384': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth',
+        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
+    'convnext_xlarge.fb_in22k_ft_in1k_384': _cfg(
+        url='https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth',
+        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
+
+    'convnext_tiny_in22k.fb_in22k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth", num_classes=21841),
+    'convnext_small_in22k.fb_in22k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth", num_classes=21841),
+    'convnext_base_in22k.fb_in22k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth", num_classes=21841),
+    'convnext_large_in22k.fb_in22k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth", num_classes=21841),
+    'convnext_xlarge_in22k.fb_in22k': _cfg(
+        url="https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth", num_classes=21841),
+})
+
+
 @register_model
 def convnext_atto(pretrained=False, **kwargs):
     # timm femto variant (NOTE: still tweaking depths, will vary between 3-4M param, current is 3.7M
@@ -569,105 +572,7 @@ def convnext_large(pretrained=False, **kwargs):
 
 
 @register_model
-def convnext_tiny_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768), **kwargs)
-    model = _create_convnext('convnext_tiny_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_small_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
-    model = _create_convnext('convnext_small_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_base_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
-    model = _create_convnext('convnext_base_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_large_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
-    model = _create_convnext('convnext_large_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_xlarge_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
-    model = _create_convnext('convnext_xlarge_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_tiny_384_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768), **kwargs)
-    model = _create_convnext('convnext_tiny_384_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_small_384_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
-    model = _create_convnext('convnext_small_384_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_base_384_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
-    model = _create_convnext('convnext_base_384_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_large_384_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
-    model = _create_convnext('convnext_large_384_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_xlarge_384_in22ft1k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
-    model = _create_convnext('convnext_xlarge_384_in22ft1k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_tiny_in22k(pretrained=False, **kwargs):
-    model_args = dict(depths=(3, 3, 9, 3), dims=(96, 192, 384, 768), **kwargs)
-    model = _create_convnext('convnext_tiny_in22k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_small_in22k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
-    model = _create_convnext('convnext_small_in22k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_base_in22k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
-    model = _create_convnext('convnext_base_in22k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_large_in22k(pretrained=False, **kwargs):
-    model_args = dict(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
-    model = _create_convnext('convnext_large_in22k', pretrained=pretrained, **model_args)
-    return model
-
-
-@register_model
-def convnext_xlarge_in22k(pretrained=False, **kwargs):
+def convnext_xlarge(pretrained=False, **kwargs):
     model_args = dict(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
-    model = _create_convnext('convnext_xlarge_in22k', pretrained=pretrained, **model_args)
+    model = _create_convnext('convnext_xlarge', pretrained=pretrained, **model_args)
     return model

timm/models/efficientnet.py

@@ -366,11 +366,11 @@ default_cfgs = {
     'tf_efficientnetv2_m': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_m-cc09e0cd.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
-        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
     'tf_efficientnetv2_l': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_l-d664b728.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
-        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
 
     'tf_efficientnetv2_s_in21ft1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_s_21ft1k-d7dafa41.pth',
@@ -379,15 +379,15 @@ default_cfgs = {
     'tf_efficientnetv2_m_in21ft1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_m_21ft1k-bf41664a.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
-        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
     'tf_efficientnetv2_l_in21ft1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_l_21ft1k-60127a9d.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
-        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
     'tf_efficientnetv2_xl_in21ft1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_xl_in21ft1k-06c35c48.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
-        input_size=(3, 384, 384), test_input_size=(3, 512, 512), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 512, 512), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
 
     'tf_efficientnetv2_s_in21k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_s_21k-6337ad01.pth',
@@ -396,15 +396,15 @@ default_cfgs = {
     'tf_efficientnetv2_m_in21k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_m_21k-361418a2.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), num_classes=21843,
-        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
     'tf_efficientnetv2_l_in21k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_l_21k-91a19ec9.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), num_classes=21843,
-        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
     'tf_efficientnetv2_xl_in21k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_xl_in21k-fd7e8abf.pth',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), num_classes=21843,
-        input_size=(3, 384, 384), test_input_size=(3, 512, 512), pool_size=(12, 12), crop_pct=1.0),
+        input_size=(3, 384, 384), test_input_size=(3, 512, 512), pool_size=(12, 12), crop_pct=1.0, crop_mode='squash'),
 
     'tf_efficientnetv2_b0': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/tf_efficientnetv2_b0-c7cc451f.pth',

timm/models/factory.py

@@ -1,14 +1,18 @@
-from urllib.parse import urlsplit, urlunsplit
 import os
+from typing import Any, Dict, Optional, Union
+from urllib.parse import urlsplit
 
-from .registry import is_model, is_model_in_modules, model_entrypoint
+from .pretrained import PretrainedCfg, split_model_name_tag
 from .helpers import load_checkpoint
-from .layers import set_layer_config
 from .hub import load_model_config_from_hf
+from .layers import set_layer_config
+from .registry import is_model, model_entrypoint
 
 
 def parse_model_name(model_name):
-    model_name = model_name.replace('hf_hub', 'hf-hub')  # NOTE for backwards compat, to deprecate hf_hub use
+    if model_name.startswith('hf_hub'):
+        # NOTE for backwards compat, deprecate hf_hub use
+        model_name = model_name.replace('hf_hub', 'hf-hub')
     parsed = urlsplit(model_name)
     assert parsed.scheme in ('', 'timm', 'hf-hub')
     if parsed.scheme == 'hf-hub':
@@ -20,6 +24,7 @@ def parse_model_name(model_name):
 
 
 def safe_model_name(model_name, remove_source=True):
+    # return a filename / path safe model name
     def make_safe(name):
         return ''.join(c if c.isalnum() else '_' for c in name).rstrip('_')
     if remove_source:
@@ -28,20 +33,29 @@ def safe_model_name(model_name, remove_source=True):
 
 
 def create_model(
-        model_name,
-        pretrained=False,
-        pretrained_cfg=None,
-        checkpoint_path='',
-        scriptable=None,
-        exportable=None,
-        no_jit=None,
-        **kwargs):
+        model_name: str,
+        pretrained: bool = False,
+        pretrained_cfg: Optional[Union[str, Dict[str, Any], PretrainedCfg]] = None,
+        pretrained_cfg_overlay:  Optional[Dict[str, Any]] = None,
+        checkpoint_path: str = '',
+        scriptable: Optional[bool] = None,
+        exportable: Optional[bool] = None,
+        no_jit: Optional[bool] = None,
+        **kwargs,
+):
     """Create a model
 
+    Lookup model's entrypoint function and pass relevant args to create a new model.
+
+    **kwargs will be passed through entrypoint fn to timm.models.build_model_with_cfg()
+    and then the model class __init__(). kwargs values set to None are pruned before passing.
+
     Args:
         model_name (str): name of model to instantiate
         pretrained (bool): load pretrained ImageNet-1k weights if true
-        checkpoint_path (str): path of checkpoint to load after model is initialized
+        pretrained_cfg (Union[str, dict, PretrainedCfg]): pass in external pretrained_cfg for model
+        pretrained_cfg_overlay (dict): replace key-values in base pretrained_cfg with these
+        checkpoint_path (str): path of checkpoint to load _after_ the model is initialized
         scriptable (bool): set layer config so that model is jit scriptable (not working for all models yet)
         exportable (bool): set layer config so that model is traceable / ONNX exportable (not fully impl/obeyed yet)
         no_jit (bool): set layer config so that model doesn't utilize jit scripted layers (so far activations only)
@@ -49,7 +63,7 @@ def create_model(
     Keyword Args:
         drop_rate (float): dropout rate for training (default: 0.0)
         global_pool (str): global pool type (default: 'avg')
-        **: other kwargs are model specific
+        **: other kwargs are consumed by builder or model __init__()
     """
     # Parameters that aren't supported by all models or are intended to only override model defaults if set
     # should default to None in command line args/cfg. Remove them if they are present and not set so that
@@ -58,17 +72,27 @@ def create_model(
 
     model_source, model_name = parse_model_name(model_name)
     if model_source == 'hf-hub':
-        # FIXME hf-hub source overrides any passed in pretrained_cfg, warn?
+        assert not pretrained_cfg, 'pretrained_cfg should not be set when sourcing model from Hugging Face Hub.'
         # For model names specified in the form `hf-hub:path/architecture_name@revision`,
         # load model weights + pretrained_cfg from Hugging Face hub.
         pretrained_cfg, model_name = load_model_config_from_hf(model_name)
+    else:
+        model_name, pretrained_tag = split_model_name_tag(model_name)
+        if not pretrained_cfg:
+            # a valid pretrained_cfg argument takes priority over tag in model name
+            pretrained_cfg = pretrained_tag
 
     if not is_model(model_name):
         raise RuntimeError('Unknown model (%s)' % model_name)
 
     create_fn = model_entrypoint(model_name)
     with set_layer_config(scriptable=scriptable, exportable=exportable, no_jit=no_jit):
-        model = create_fn(pretrained=pretrained, pretrained_cfg=pretrained_cfg, **kwargs)
+        model = create_fn(
+            pretrained=pretrained,
+            pretrained_cfg=pretrained_cfg,
+            pretrained_cfg_overlay=pretrained_cfg_overlay,
+            **kwargs,
+        )
 
     if checkpoint_path:
         load_checkpoint(model, checkpoint_path)

timm/models/helpers.py

@@ -3,6 +3,7 @@
 Hacked together by / Copyright 2020 Ross Wightman
 """
 import collections.abc
+import dataclasses
 import logging
 import math
 import os
@@ -17,6 +18,7 @@ import torch.nn as nn
 from torch.hub import load_state_dict_from_url
 from torch.utils.checkpoint import checkpoint
 
+from .pretrained import PretrainedCfg
 from .features import FeatureListNet, FeatureDictNet, FeatureHookNet
 from .fx_features import FeatureGraphNet
 from .hub import has_hf_hub, download_cached_file, load_state_dict_from_hf
@@ -156,7 +158,7 @@ def _resolve_pretrained_source(pretrained_cfg):
             # hf-hub available as alternate weight source in default_cfg
             load_from = 'hf-hub'
             pretrained_loc = hf_hub_id
-    if load_from == 'hf-hub' and 'hf_hub_filename' in pretrained_cfg:
+    if load_from == 'hf-hub' and pretrained_cfg.get('hf_hub_filename', None):
         # if a filename override is set, return tuple for location w/ (hub_id, filename)
         pretrained_loc = pretrained_loc, pretrained_cfg['hf_hub_filename']
     return load_from, pretrained_loc
@@ -191,10 +193,14 @@ def load_custom_pretrained(
     Args:
         model: The instantiated model to load weights into
         pretrained_cfg (dict): Default pretrained model cfg
-        load_fn: An external stand alone fn that loads weights into provided model, otherwise a fn named
+        load_fn: An external standalone fn that loads weights into provided model, otherwise a fn named
             'laod_pretrained' on the model will be called if it exists
     """
-    pretrained_cfg = pretrained_cfg or getattr(model, 'pretrained_cfg', None) or {}
+    pretrained_cfg = pretrained_cfg or getattr(model, 'pretrained_cfg', None)
+    if not pretrained_cfg:
+        _logger.warning("Invalid pretrained config, cannot load weights.")
+        return
+
     load_from, pretrained_loc = _resolve_pretrained_source(pretrained_cfg)
     if not load_from:
         _logger.warning("No pretrained weights exist for this model. Using random initialization.")
@@ -202,7 +208,11 @@ def load_custom_pretrained(
     if load_from == 'hf-hub':  # FIXME
         _logger.warning("Hugging Face hub not currently supported for custom load pretrained models.")
     elif load_from == 'url':
-        pretrained_loc = download_cached_file(pretrained_loc, check_hash=_CHECK_HASH, progress=_DOWNLOAD_PROGRESS)
+        pretrained_loc = download_cached_file(
+            pretrained_loc,
+            check_hash=_CHECK_HASH,
+            progress=_DOWNLOAD_PROGRESS
+        )
 
     if load_fn is not None:
         load_fn(model, pretrained_loc)
@@ -250,13 +260,17 @@ def load_pretrained(
     Args:
         model (nn.Module) : PyTorch model module
         pretrained_cfg (Optional[Dict]): configuration for pretrained weights / target dataset
-        num_classes (int): num_classes for model
-        in_chans (int): in_chans for model
+        num_classes (int): num_classes for target model
+        in_chans (int): in_chans for target model
         filter_fn (Optional[Callable]): state_dict filter fn for load (takes state_dict, model as args)
         strict (bool): strict load of checkpoint
 
     """
-    pretrained_cfg = pretrained_cfg or getattr(model, 'pretrained_cfg', None) or {}
+    pretrained_cfg = pretrained_cfg or getattr(model, 'pretrained_cfg', None)
+    if not pretrained_cfg:
+        _logger.warning("Invalid pretrained config, cannot load weights.")
+        return
+
     load_from, pretrained_loc = _resolve_pretrained_source(pretrained_cfg)
     if load_from == 'file':
         _logger.info(f'Loading pretrained weights from file ({pretrained_loc})')
@@ -264,7 +278,11 @@ def load_pretrained(
     elif load_from == 'url':
         _logger.info(f'Loading pretrained weights from url ({pretrained_loc})')
         state_dict = load_state_dict_from_url(
-            pretrained_loc, map_location='cpu', progress=_DOWNLOAD_PROGRESS, check_hash=_CHECK_HASH)
+            pretrained_loc,
+            map_location='cpu',
+            progress=_DOWNLOAD_PROGRESS,
+            check_hash=_CHECK_HASH,
+        )
     elif load_from == 'hf-hub':
         _logger.info(f'Loading pretrained weights from Hugging Face hub ({pretrained_loc})')
         if isinstance(pretrained_loc, (list, tuple)):
@@ -428,40 +446,20 @@ def adapt_model_from_file(parent_module, model_variant):
 def pretrained_cfg_for_features(pretrained_cfg):
     pretrained_cfg = deepcopy(pretrained_cfg)
     # remove default pretrained cfg fields that don't have much relevance for feature backbone
-    to_remove = ('num_classes', 'crop_pct', 'classifier', 'global_pool')  # add default final pool size?
+    to_remove = ('num_classes', 'classifier', 'global_pool')  # add default final pool size?
     for tr in to_remove:
         pretrained_cfg.pop(tr, None)
     return pretrained_cfg
 
 
-def set_default_kwargs(kwargs, names, pretrained_cfg):
-    for n in names:
-        # for legacy reasons, model __init__args uses img_size + in_chans as separate args while
-        # pretrained_cfg has one input_size=(C, H ,W) entry
-        if n == 'img_size':
-            input_size = pretrained_cfg.get('input_size', None)
-            if input_size is not None:
-                assert len(input_size) == 3
-                kwargs.setdefault(n, input_size[-2:])
-        elif n == 'in_chans':
-            input_size = pretrained_cfg.get('input_size', None)
-            if input_size is not None:
-                assert len(input_size) == 3
-                kwargs.setdefault(n, input_size[0])
-        else:
-            default_val = pretrained_cfg.get(n, None)
-            if default_val is not None:
-                kwargs.setdefault(n, pretrained_cfg[n])
-
-
-def filter_kwargs(kwargs, names):
+def _filter_kwargs(kwargs, names):
     if not kwargs or not names:
         return
     for n in names:
         kwargs.pop(n, None)
 
 
-def update_pretrained_cfg_and_kwargs(pretrained_cfg, kwargs, kwargs_filter):
+def _update_default_kwargs(pretrained_cfg, kwargs, kwargs_filter):
     """ Update the default_cfg and kwargs before passing to model
 
     Args:
@@ -474,31 +472,61 @@ def update_pretrained_cfg_and_kwargs(pretrained_cfg, kwargs, kwargs_filter):
     if pretrained_cfg.get('fixed_input_size', False):
         # if fixed_input_size exists and is True, model takes an img_size arg that fixes its input size
         default_kwarg_names += ('img_size',)
-    set_default_kwargs(kwargs, names=default_kwarg_names, pretrained_cfg=pretrained_cfg)
+
+    for n in default_kwarg_names:
+        # for legacy reasons, model __init__args uses img_size + in_chans as separate args while
+        # pretrained_cfg has one input_size=(C, H ,W) entry
+        if n == 'img_size':
+            input_size = pretrained_cfg.get('input_size', None)
+            if input_size is not None:
+                assert len(input_size) == 3
+                kwargs.setdefault(n, input_size[-2:])
+        elif n == 'in_chans':
+            input_size = pretrained_cfg.get('input_size', None)
+            if input_size is not None:
+                assert len(input_size) == 3
+                kwargs.setdefault(n, input_size[0])
+        else:
+            default_val = pretrained_cfg.get(n, None)
+            if default_val is not None:
+                kwargs.setdefault(n, pretrained_cfg[n])
+
     # Filter keyword args for task specific model variants (some 'features only' models, etc.)
-    filter_kwargs(kwargs, names=kwargs_filter)
+    _filter_kwargs(kwargs, names=kwargs_filter)
 
 
-def resolve_pretrained_cfg(variant: str, pretrained_cfg=None):
-    if pretrained_cfg and isinstance(pretrained_cfg, dict):
-        # highest priority, pretrained_cfg available and passed as arg
-        return deepcopy(pretrained_cfg)
+def resolve_pretrained_cfg(
+        variant: str,
+        pretrained_cfg=None,
+        pretrained_cfg_overlay=None,
+) -> PretrainedCfg:
+    model_with_tag = variant
+    pretrained_tag = None
+    if pretrained_cfg:
+        if isinstance(pretrained_cfg, dict):
+            # pretrained_cfg dict passed as arg, validate by converting to PretrainedCfg
+            pretrained_cfg = PretrainedCfg(**pretrained_cfg)
+        elif isinstance(pretrained_cfg, str):
+            pretrained_tag = pretrained_cfg
+            pretrained_cfg = None
+
     # fallback to looking up pretrained cfg in model registry by variant identifier
-    pretrained_cfg = get_pretrained_cfg(variant)
+    if not pretrained_cfg:
+        if pretrained_tag:
+            model_with_tag = '.'.join([variant, pretrained_tag])
+        pretrained_cfg = get_pretrained_cfg(model_with_tag)
+
     if not pretrained_cfg:
         _logger.warning(
-            f"No pretrained configuration specified for {variant} model. Using a default."
+            f"No pretrained configuration specified for {model_with_tag} model. Using a default."
             f" Please add a config to the model pretrained_cfg registry or pass explicitly.")
-        pretrained_cfg = dict(
-            url='',
-            num_classes=1000,
-            input_size=(3, 224, 224),
-            pool_size=None,
-            crop_pct=.9,
-            interpolation='bicubic',
-            first_conv='',
-            classifier='',
-        )
+        pretrained_cfg = PretrainedCfg()  # instance with defaults
+
+    pretrained_cfg_overlay = pretrained_cfg_overlay or {}
+    if not pretrained_cfg.architecture:
+        pretrained_cfg_overlay.setdefault('architecture', variant)
+    pretrained_cfg = dataclasses.replace(pretrained_cfg, **pretrained_cfg_overlay)
+
     return pretrained_cfg
 
 
@@ -507,13 +535,14 @@ def build_model_with_cfg(
         variant: str,
         pretrained: bool,
         pretrained_cfg: Optional[Dict] = None,
+        pretrained_cfg_overlay: Optional[Dict] = None,
         model_cfg: Optional[Any] = None,
         feature_cfg: Optional[Dict] = None,
         pretrained_strict: bool = True,
         pretrained_filter_fn: Optional[Callable] = None,
-        pretrained_custom_load: bool = False,
         kwargs_filter: Optional[Tuple[str]] = None,
-        **kwargs):
+        **kwargs,
+):
     """ Build model with specified default_cfg and optional model_cfg
 
     This helper fn aids in the construction of a model including:
@@ -531,7 +560,6 @@ def build_model_with_cfg(
         feature_cfg (Optional[Dict]: feature extraction adapter config
         pretrained_strict (bool): load pretrained weights strictly
         pretrained_filter_fn (Optional[Callable]): filter callable for pretrained weights
-        pretrained_custom_load (bool): use custom load fn, to load numpy or other non PyTorch weights
         kwargs_filter (Optional[Tuple]): kwargs to filter before passing to model
         **kwargs: model args passed through to model __init__
     """
@@ -540,9 +568,16 @@ def build_model_with_cfg(
     feature_cfg = feature_cfg or {}
 
     # resolve and update model pretrained config and model kwargs
-    pretrained_cfg = resolve_pretrained_cfg(variant, pretrained_cfg=pretrained_cfg)
-    update_pretrained_cfg_and_kwargs(pretrained_cfg, kwargs, kwargs_filter)
-    pretrained_cfg.setdefault('architecture', variant)
+    pretrained_cfg = resolve_pretrained_cfg(
+        variant,
+        pretrained_cfg=pretrained_cfg,
+        pretrained_cfg_overlay=pretrained_cfg_overlay
+    )
+
+    # FIXME converting back to dict, PretrainedCfg use should be propagated further, but not into model
+    pretrained_cfg = pretrained_cfg.to_dict()
+
+    _update_default_kwargs(pretrained_cfg, kwargs, kwargs_filter)
 
     # Setup for feature extraction wrapper done at end of this fn
     if kwargs.pop('features_only', False):
@@ -551,8 +586,11 @@ def build_model_with_cfg(
         if 'out_indices' in kwargs:
             feature_cfg['out_indices'] = kwargs.pop('out_indices')
 
-    # Build the model
-    model = model_cls(**kwargs) if model_cfg is None else model_cls(cfg=model_cfg, **kwargs)
+    # Instantiate the model
+    if model_cfg is None:
+        model = model_cls(**kwargs)
+    else:
+        model = model_cls(cfg=model_cfg, **kwargs)
     model.pretrained_cfg = pretrained_cfg
     model.default_cfg = model.pretrained_cfg  # alias for backwards compat
     
@@ -562,9 +600,11 @@ def build_model_with_cfg(
     # For classification models, check class attr, then kwargs, then default to 1k, otherwise 0 for feats
     num_classes_pretrained = 0 if features else getattr(model, 'num_classes', kwargs.get('num_classes', 1000))
     if pretrained:
-        if pretrained_custom_load:
-            # FIXME improve custom load trigger
-            load_custom_pretrained(model, pretrained_cfg=pretrained_cfg)
+        if pretrained_cfg.get('custom_load', False):
+            load_custom_pretrained(
+                model,
+                pretrained_cfg=pretrained_cfg,
+            )
         else:
             load_pretrained(
                 model,
@@ -572,7 +612,8 @@ def build_model_with_cfg(
                 num_classes=num_classes_pretrained,
                 in_chans=kwargs.get('in_chans', 3),
                 filter_fn=pretrained_filter_fn,
-                strict=pretrained_strict)
+                strict=pretrained_strict,
+            )
 
     # Wrap the model in a feature extraction module if enabled
     if features:

timm/models/hub.py

@@ -15,11 +15,13 @@ except ImportError:
     from torch.hub import _get_torch_home as get_dir
 
 from timm import __version__
+from timm.models.pretrained import filter_pretrained_cfg
 
 try:
-    from huggingface_hub import (create_repo, get_hf_file_metadata,
-                                 hf_hub_download, hf_hub_url,
-                                 repo_type_and_id_from_hf_id, upload_folder)
+    from huggingface_hub import (
+        create_repo, get_hf_file_metadata,
+        hf_hub_download, hf_hub_url,
+        repo_type_and_id_from_hf_id, upload_folder)
     from huggingface_hub.utils import EntryNotFoundError
     hf_hub_download = partial(hf_hub_download, library_name="timm", library_version=__version__)
     _has_hf_hub = True
@@ -46,8 +48,11 @@ def get_cache_dir(child_dir=''):
 
 
 def download_cached_file(url, check_hash=True, progress=False):
-    parts = urlparse(url)
-    filename = os.path.basename(parts.path)
+    if isinstance(url, (list, tuple)):
+        url, filename = url
+    else:
+        parts = urlparse(url)
+        filename = os.path.basename(parts.path)
     cached_file = os.path.join(get_cache_dir(), filename)
     if not os.path.exists(cached_file):
         _logger.info('Downloading: "{}" to {}\n'.format(url, cached_file))
@@ -90,10 +95,27 @@ def _download_from_hf(model_id: str, filename: str):
 def load_model_config_from_hf(model_id: str):
     assert has_hf_hub(True)
     cached_file = _download_from_hf(model_id, 'config.json')
-    pretrained_cfg = load_cfg_from_json(cached_file)
+
+    hf_config = load_cfg_from_json(cached_file)
+    if 'pretrained_cfg' not in hf_config:
+        # old form, pull pretrain_cfg out of the base dict
+        pretrained_cfg = hf_config
+        hf_config = {}
+        hf_config['architecture'] = pretrained_cfg.pop('architecture')
+        hf_config['num_features'] = pretrained_cfg.pop('num_features', None)
+        if 'labels' in pretrained_cfg:
+            hf_config['label_name'] = pretrained_cfg.pop('labels')
+        hf_config['pretrained_cfg'] = pretrained_cfg
+
+    # NOTE currently discarding parent config as only arch name and pretrained_cfg used in timm right now
+    pretrained_cfg = hf_config['pretrained_cfg']
     pretrained_cfg['hf_hub_id'] = model_id  # insert hf_hub id for pretrained weight load during model creation
     pretrained_cfg['source'] = 'hf-hub'
-    model_name = pretrained_cfg.get('architecture')
+    if 'num_classes' in hf_config:
+        # model should be created with parent num_classes if they exist
+        pretrained_cfg['num_classes'] = hf_config['num_classes']
+    model_name = hf_config['architecture']
+
     return pretrained_cfg, model_name
 
 
@@ -114,10 +136,34 @@ def save_for_hf(model, save_directory, model_config=None):
     torch.save(model.state_dict(), weights_path)
 
     config_path = save_directory / 'config.json'
-    hf_config = model.pretrained_cfg
-    hf_config['num_classes'] = model_config.pop('num_classes', model.num_classes)
-    hf_config['num_features'] = model_config.pop('num_features', model.num_features)
-    hf_config['labels'] = model_config.pop('labels', [f"LABEL_{i}" for i in range(hf_config['num_classes'])])
+    hf_config = {}
+    pretrained_cfg = filter_pretrained_cfg(model.pretrained_cfg, remove_source=True, remove_null=True)
+    # set some values at root config level
+    hf_config['architecture'] = pretrained_cfg.pop('architecture')
+    hf_config['num_classes'] = model_config.get('num_classes', model.num_classes)
+    hf_config['num_features'] = model_config.get('num_features', model.num_features)
+    hf_config['global_pool'] = model_config.get('global_pool', getattr(model, 'global_pool', None))
+
+    if 'label' in model_config:
+        _logger.warning(
+            "'label' as a config field for timm models is deprecated. Please use 'label_name' and 'display_name'. "
+            "Using provided 'label' field as 'label_name'.")
+        model_config['label_name'] = model_config.pop('label')
+
+    label_name = model_config.pop('label_name', None)
+    if label_name:
+        assert isinstance(label_name, (dict, list, tuple))
+        # map label id (classifier index) -> unique label name (ie synset for ImageNet, MID for OpenImages)
+        # can be a dict id: name if there are id gaps, or tuple/list if no gaps.
+        hf_config['label_name'] = model_config['label_name']
+
+    display_name = model_config.pop('display_name', None)
+    if display_name:
+        assert isinstance(display_name, dict)
+        # map label_name -> user interface display name
+        hf_config['display_name'] = model_config['display_name']
+
+    hf_config['pretrained_cfg'] = pretrained_cfg
     hf_config.update(model_config)
 
     with config_path.open('w') as f:
@@ -127,14 +173,14 @@ def save_for_hf(model, save_directory, model_config=None):
 def push_to_hf_hub(
     model,
     repo_id: str,
-    commit_message: str ='Add model',
+    commit_message: str = 'Add model',
     token: Optional[str] = None,
     revision: Optional[str] = None,
     private: bool = False,
     create_pr: bool = False,
     model_config: Optional[dict] = None,
 ):
-    # Create repo if doesn't exist yet
+    # Create repo if it doesn't exist yet
     repo_url = create_repo(repo_id, token=token, private=private, exist_ok=True)
 
     # Infer complete repo_id from repo_url
@@ -154,10 +200,11 @@ def push_to_hf_hub(
         # Save model weights and config.
         save_for_hf(model, tmpdir, model_config=model_config)
 
-        # Add readme if does not exist
+        # Add readme if it does not exist
         if not has_readme:
+            model_name = repo_id.split('/')[-1]
             readme_path = Path(tmpdir) / "README.md"
-            readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_id}'
+            readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {model_name}'
             readme_path.write_text(readme_text)
 
         # Upload model and return

timm/models/inception_v3.py

@@ -27,24 +27,23 @@ def _cfg(url='', **kwargs):
 default_cfgs = {
     # original PyTorch weights, ported from Tensorflow but modified
     'inception_v3': _cfg(
-        url='https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth',
-        has_aux=True),  # checkpoint has aux logit layer weights
+        # NOTE checkpoint has aux logit layer weights
+        url='https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth'),
     # my port of Tensorflow SLIM weights (http://download.tensorflow.org/models/inception_v3_2016_08_28.tar.gz)
     'tf_inception_v3': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_inception_v3-e0069de4.pth',
-        num_classes=1000, has_aux=False, label_offset=1),
+        num_classes=1000, label_offset=1),
     # my port of Tensorflow adversarially trained Inception V3 from
     # http://download.tensorflow.org/models/adv_inception_v3_2017_08_18.tar.gz
     'adv_inception_v3': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/adv_inception_v3-9e27bd63.pth',
-        num_classes=1000, has_aux=False, label_offset=1),
+        num_classes=1000, label_offset=1),
     # from gluon pretrained models, best performing in terms of accuracy/loss metrics
     # https://gluon-cv.mxnet.io/model_zoo/classification.html
     'gluon_inception_v3': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/gluon_inception_v3-9f746940.pth',
         mean=IMAGENET_DEFAULT_MEAN,  # also works well with inception defaults
         std=IMAGENET_DEFAULT_STD,  # also works well with inception defaults
-        has_aux=False,
     )
 }
 
@@ -433,10 +432,10 @@ def _create_inception_v3(variant, pretrained=False, **kwargs):
     if aux_logits:
         assert not kwargs.pop('features_only', False)
         model_cls = InceptionV3Aux
-        load_strict = pretrained_cfg['has_aux']
+        load_strict = variant == 'inception_v3'
     else:
         model_cls = InceptionV3
-        load_strict = not pretrained_cfg['has_aux']
+        load_strict = variant != 'inception_v3'
 
     return build_model_with_cfg(
         model_cls, variant, pretrained,

timm/models/layers/activations.py

@@ -143,3 +143,17 @@ class GELU(nn.Module):
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         return F.gelu(input)
+
+
+def gelu_tanh(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return F.gelu(x, approximate='tanh')
+
+
+class GELUTanh(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(GELUTanh, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.gelu(input, approximate='tanh')

timm/models/layers/create_act.py

@@ -28,6 +28,7 @@ _ACT_FN_DEFAULT = dict(
     celu=F.celu,
     selu=F.selu,
     gelu=gelu,
+    gelu_tanh=gelu_tanh,
     sigmoid=sigmoid,
     tanh=tanh,
     hard_sigmoid=F.hardsigmoid if _has_hardsigmoid else hard_sigmoid,
@@ -71,6 +72,7 @@ _ACT_LAYER_DEFAULT = dict(
     celu=nn.CELU,
     selu=nn.SELU,
     gelu=GELU,
+    gelu_tanh=GELUTanh,
     sigmoid=Sigmoid,
     tanh=Tanh,
     hard_sigmoid=nn.Hardsigmoid if _has_hardsigmoid else HardSigmoid,

timm/models/pretrained.py

@@ -0,0 +1,122 @@
+import copy
+from collections import deque, defaultdict
+from dataclasses import dataclass, field, replace, asdict
+from typing import Any, Deque, Dict, Tuple, Optional, Union
+
+
+@dataclass
+class PretrainedCfg:
+    """
+    """
+    # weight locations
+    url: Optional[Union[str, Tuple[str, str]]] = None
+    file: Optional[str] = None
+    hf_hub_id: Optional[str] = None
+    hf_hub_filename: Optional[str] = None
+
+    source: Optional[str] = None  # source of cfg / weight location used (url, file, hf-hub)
+    architecture: Optional[str] = None  # architecture variant can be set when not implicit
+    custom_load: bool = False  # use custom model specific model.load_pretrained() (ie for npz files)
+
+    # input / data config
+    input_size: Tuple[int, int, int] = (3, 224, 224)
+    test_input_size: Optional[Tuple[int, int, int]] = None
+    min_input_size: Optional[Tuple[int, int, int]] = None
+    fixed_input_size: bool = False
+    interpolation: str = 'bicubic'
+    crop_pct: float = 0.875
+    test_crop_pct: Optional[float] = None
+    crop_mode: str = 'center'
+    mean: Tuple[float, ...] = (0.485, 0.456, 0.406)
+    std: Tuple[float, ...] = (0.229, 0.224, 0.225)
+
+    # head config
+    num_classes: int = 1000
+    label_offset: Optional[int] = None
+
+    # model attributes that vary with above or required for pretrained adaptation
+    pool_size: Optional[Tuple[int, ...]] = None
+    test_pool_size: Optional[Tuple[int, ...]] = None
+    first_conv: Optional[str] = None
+    classifier: Optional[str] = None
+
+    license: Optional[str] = None
+    source_url: Optional[str] = None
+    paper: Optional[str] = None
+    notes: Optional[str] = None
+
+    @property
+    def has_weights(self):
+        return self.url or self.file or self.hf_hub_id
+
+    def to_dict(self, remove_source=False, remove_null=True):
+        return filter_pretrained_cfg(
+            asdict(self),
+            remove_source=remove_source,
+            remove_null=remove_null
+        )
+
+
+def filter_pretrained_cfg(cfg, remove_source=False, remove_null=True):
+    filtered_cfg = {}
+    for k, v in cfg.items():
+        if remove_source and k in {'url', 'file', 'hf_hub_id', 'hf_hub_id', 'hf_hub_filename', 'source'}:
+            continue
+        if remove_null and v is None:
+            continue
+        filtered_cfg[k] = v
+    return filtered_cfg
+
+
+@dataclass
+class DefaultCfg:
+    tags: Deque[str] = field(default_factory=deque)  # priority queue of tags (first is default)
+    cfgs: Dict[str, PretrainedCfg] = field(default_factory=dict)  # pretrained cfgs by tag
+    is_pretrained: bool = False  # at least one of the configs has a pretrained source set
+
+    @property
+    def default(self):
+        return self.cfgs[self.tags[0]]
+
+    @property
+    def default_with_tag(self):
+        tag = self.tags[0]
+        return tag, self.cfgs[tag]
+
+
+def split_model_name_tag(model_name: str, no_tag=''):
+    model_name, *tag_list = model_name.split('.', 1)
+    tag = tag_list[0] if tag_list else no_tag
+    return model_name, tag
+
+
+def generate_default_cfgs(cfgs: Dict[str, Union[Dict[str, Any], PretrainedCfg]]):
+    out = defaultdict(DefaultCfg)
+    default_set = set()  # no tag and tags ending with * are prioritized as default
+
+    for k, v in cfgs.items():
+        if isinstance(v, dict):
+            v = PretrainedCfg(**v)
+        has_weights = v.has_weights
+
+        model, tag = split_model_name_tag(k)
+        is_default_set = model in default_set
+        priority = (has_weights and not tag) or (tag.endswith('*') and not is_default_set)
+        tag = tag.strip('*')
+
+        default_cfg = out[model]
+
+        if priority:
+            default_cfg.tags.appendleft(tag)
+            default_set.add(model)
+        elif has_weights and not default_cfg.is_pretrained:
+            default_cfg.tags.appendleft(tag)
+        else:
+            default_cfg.tags.append(tag)
+
+        if has_weights:
+            default_cfg.is_pretrained = True
+
+        default_cfg.cfgs[tag] = v
+
+    return out

timm/models/registry.py

@@ -2,20 +2,30 @@
 Hacked together by / Copyright 2020 Ross Wightman
 """
 
-import sys
-import re
 import fnmatch
-from collections import defaultdict
+import re
+import sys
+from collections import defaultdict, deque
 from copy import deepcopy
+from typing import List, Optional, Union, Tuple
+
+from .pretrained import PretrainedCfg, DefaultCfg, split_model_name_tag
 
-__all__ = ['list_models', 'is_model', 'model_entrypoint', 'list_modules', 'is_model_in_modules',
-           'is_pretrained_cfg_key', 'has_pretrained_cfg_key', 'get_pretrained_cfg_value', 'is_model_pretrained']
+__all__ = [
+    'list_models', 'is_model', 'model_entrypoint', 'list_modules', 'is_model_in_modules',
+    'get_pretrained_cfg_value', 'is_model_pretrained', 'get_arch_name']
 
 _module_to_models = defaultdict(set)  # dict of sets to check membership of model in module
 _model_to_module = {}  # mapping of model names to module names
-_model_entrypoints = {}  # mapping of model names to entrypoint fns
+_model_entrypoints = {}  # mapping of model names to architecture entrypoint fns
 _model_has_pretrained = set()  # set of model names that have pretrained weight url present
-_model_pretrained_cfgs = dict()  # central repo for model default_cfgs
+_model_default_cfgs = dict()  # central repo for model arch -> default cfg objects
+_model_pretrained_cfgs = dict()  # central repo for model arch + tag -> pretrained cfgs
+_model_with_tags = defaultdict(list)  # shortcut to map each model arch to all model + tag names
+
+
+def get_arch_name(model_name: str) -> Tuple[str, Optional[str]]:
+    return split_model_name_tag(model_name)[0]
 
 
 def register_model(fn):
@@ -35,19 +45,37 @@ def register_model(fn):
     _model_entrypoints[model_name] = fn
     _model_to_module[model_name] = module_name
     _module_to_models[module_name].add(model_name)
-    has_valid_pretrained = False  # check if model has a pretrained url to allow filtering on this
     if hasattr(mod, 'default_cfgs') and model_name in mod.default_cfgs:
         # this will catch all models that have entrypoint matching cfg key, but miss any aliasing
         # entrypoints or non-matching combos
         cfg = mod.default_cfgs[model_name]
-        has_valid_pretrained = (
-            ('url' in cfg and 'http' in cfg['url']) or
-            ('file' in cfg and cfg['file']) or
-            ('hf_hub_id' in cfg and cfg['hf_hub_id'])
-        )
-        _model_pretrained_cfgs[model_name] = mod.default_cfgs[model_name]
-    if has_valid_pretrained:
-        _model_has_pretrained.add(model_name)
+        if not isinstance(cfg, DefaultCfg):
+            # new style default cfg dataclass w/ multiple entries per model-arch
+            assert isinstance(cfg, dict)
+            # old style cfg dict per model-arch
+            cfg = PretrainedCfg(**cfg)
+            cfg = DefaultCfg(tags=deque(['']), cfgs={'': cfg})
+
+        for tag_idx, tag in enumerate(cfg.tags):
+            is_default = tag_idx == 0
+            pretrained_cfg = cfg.cfgs[tag]
+            if is_default:
+                _model_pretrained_cfgs[model_name] = pretrained_cfg
+                if pretrained_cfg.has_weights:
+                    # add tagless entry if it's default and has weights
+                    _model_has_pretrained.add(model_name)
+            if tag:
+                model_name_tag = '.'.join([model_name, tag])
+                _model_pretrained_cfgs[model_name_tag] = pretrained_cfg
+                if pretrained_cfg.has_weights:
+                    # add model w/ tag if tag is valid
+                    _model_has_pretrained.add(model_name_tag)
+                _model_with_tags[model_name].append(model_name_tag)
+            else:
+                _model_with_tags[model_name].append(model_name)  # has empty tag (to slowly remove these instances)
+
+        _model_default_cfgs[model_name] = cfg
+
     return fn
 
 
@@ -55,24 +83,44 @@ def _natural_key(string_):
     return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
 
 
-def list_models(filter='', module='', pretrained=False, exclude_filters='', name_matches_cfg=False):
+def list_models(
+        filter: Union[str, List[str]] = '',
+        module: str = '',
+        pretrained=False,
+        exclude_filters: str = '',
+        name_matches_cfg: bool = False,
+        include_tags: Optional[bool] = None,
+):
     """ Return list of available model names, sorted alphabetically
 
     Args:
         filter (str) - Wildcard filter string that works with fnmatch
-        module (str) - Limit model selection to a specific sub-module (ie 'gen_efficientnet')
-        pretrained (bool) - Include only models with pretrained weights if True
+        module (str) - Limit model selection to a specific submodule (ie 'vision_transformer')
+        pretrained (bool) - Include only models with valid pretrained weights if True
         exclude_filters (str or list[str]) - Wildcard filters to exclude models after including them with filter
         name_matches_cfg (bool) - Include only models w/ model_name matching default_cfg name (excludes some aliases)
-
+        include_tags (Optional[boo]) - Include pretrained tags in model names (model.tag). If None, defaults
+            set to True when pretrained=True else False (default: None)
     Example:
         model_list('gluon_resnet*') -- returns all models starting with 'gluon_resnet'
         model_list('*resnext*, 'resnet') -- returns all models with 'resnext' in 'resnet' module
     """
+    if include_tags is None:
+        # FIXME should this be default behaviour? or default to include_tags=True?
+        include_tags = pretrained
+
     if module:
         all_models = list(_module_to_models[module])
     else:
         all_models = _model_entrypoints.keys()
+
+    if include_tags:
+        # expand model names to include names w/ pretrained tags
+        models_with_tags = []
+        for m in all_models:
+            models_with_tags.extend(_model_with_tags[m])
+        all_models = models_with_tags
+
     if filter:
         models = []
         include_filters = filter if isinstance(filter, (tuple, list)) else [filter]
@@ -82,6 +130,7 @@ def list_models(filter='', module='', pretrained=False, exclude_filters='', name
                 models = set(models).union(include_models)
     else:
         models = all_models
+
     if exclude_filters:
         if not isinstance(exclude_filters, (tuple, list)):
             exclude_filters = [exclude_filters]
@@ -89,23 +138,40 @@ def list_models(filter='', module='', pretrained=False, exclude_filters='', name
             exclude_models = fnmatch.filter(models, xf)  # exclude these models
             if len(exclude_models):
                 models = set(models).difference(exclude_models)
+
     if pretrained:
         models = _model_has_pretrained.intersection(models)
+
     if name_matches_cfg:
         models = set(_model_pretrained_cfgs).intersection(models)
+
     return list(sorted(models, key=_natural_key))
 
 
+def list_pretrained(
+        filter: Union[str, List[str]] = '',
+        exclude_filters: str = '',
+):
+    return list_models(
+        filter=filter,
+        pretrained=True,
+        exclude_filters=exclude_filters,
+        include_tags=True,
+    )
+
+
 def is_model(model_name):
     """ Check if a model name exists
     """
-    return model_name in _model_entrypoints
+    arch_name = get_arch_name(model_name)
+    return arch_name in _model_entrypoints
 
 
 def model_entrypoint(model_name):
     """Fetch a model entrypoint for specified model name
     """
-    return _model_entrypoints[model_name]
+    arch_name = get_arch_name(model_name)
+    return _model_entrypoints[arch_name]
 
 
 def list_modules():
@@ -121,8 +187,9 @@ def is_model_in_modules(model_name, module_names):
         model_name (str) - name of model to check
         module_names (tuple, list, set) - names of modules to search in
     """
+    arch_name = get_arch_name(model_name)
     assert isinstance(module_names, (tuple, list, set))
-    return any(model_name in _module_to_models[n] for n in module_names)
+    return any(arch_name in _module_to_models[n] for n in module_names)
 
 
 def is_model_pretrained(model_name):
@@ -132,28 +199,12 @@ def is_model_pretrained(model_name):
 def get_pretrained_cfg(model_name):
     if model_name in _model_pretrained_cfgs:
         return deepcopy(_model_pretrained_cfgs[model_name])
-    return {}
-
-
-def has_pretrained_cfg_key(model_name, cfg_key):
-    """ Query model default_cfgs for existence of a specific key.
-    """
-    if model_name in _model_pretrained_cfgs and cfg_key in _model_pretrained_cfgs[model_name]:
-        return True
-    return False
-
-
-def is_pretrained_cfg_key(model_name, cfg_key):
-    """ Return truthy value for specified model default_cfg key, False if does not exist.
-    """
-    if model_name in _model_pretrained_cfgs and _model_pretrained_cfgs[model_name].get(cfg_key, False):
-        return True
-    return False
+    raise RuntimeError(f'No pretrained config exists for model {model_name}.')
 
 
 def get_pretrained_cfg_value(model_name, cfg_key):
-    """ Get a specific model default_cfg value by key. None if it doesn't exist.
+    """ Get a specific model default_cfg value by key. None if key doesn't exist.
     """
     if model_name in _model_pretrained_cfgs:
-        return _model_pretrained_cfgs[model_name].get(cfg_key, None)
-    return None
+        return getattr(_model_pretrained_cfgs[model_name], cfg_key, None)
+    raise RuntimeError(f'No pretrained config exist for model {model_name}.')

timm/models/regnet.py

@@ -76,6 +76,9 @@ model_cfgs = dict(
     regnety_120=RegNetCfg(w0=168, wa=73.36, wm=2.37, group_size=112, depth=19, se_ratio=0.25),
     regnety_160=RegNetCfg(w0=200, wa=106.23, wm=2.48, group_size=112, depth=18, se_ratio=0.25),
     regnety_320=RegNetCfg(w0=232, wa=115.89, wm=2.53, group_size=232, depth=20, se_ratio=0.25),
+    regnety_640=RegNetCfg(w0=352, wa=147.48, wm=2.4, group_size=328, depth=20, se_ratio=0.25),
+    regnety_1280=RegNetCfg(w0=456, wa=160.83, wm=2.52, group_size=264, depth=27, se_ratio=0.25),
+    regnety_2560=RegNetCfg(w0=640, wa=124.47, wm=2.04, group_size=848, depth=27, se_ratio=0.25),
 
     # Experimental
     regnety_040s_gn=RegNetCfg(
@@ -150,7 +153,12 @@ default_cfgs = dict(
     regnety_160=_cfg(
         url='https://dl.fbaipublicfiles.com/deit/regnety_160-a5fe301d.pth',  # from Facebook DeiT GitHub repository
         crop_pct=1.0, test_input_size=(3, 288, 288)),
-    regnety_320=_cfg(url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-regnet/regnety_320-ba464b29.pth'),
+    regnety_320=_cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-regnet/regnety_320-ba464b29.pth'
+    ),
+    regnety_640=_cfg(url=''),
+    regnety_1280=_cfg(url=''),
+    regnety_2560=_cfg(url=''),
 
     regnety_040s_gn=_cfg(url=''),
     regnetv_040=_cfg(
@@ -508,6 +516,34 @@ def _init_weights(module, name='', zero_init_last=False):
 
 def _filter_fn(state_dict):
     """ convert patch embedding weight from manual patchify + linear proj to conv"""
+    if 'classy_state_dict' in state_dict:
+        import re
+        state_dict = state_dict['classy_state_dict']['base_model']['model']
+        out = {}
+        for k, v in state_dict['trunk'].items():
+            k = k.replace('_feature_blocks.conv1.stem.0', 'stem.conv')
+            k = k.replace('_feature_blocks.conv1.stem.1', 'stem.bn')
+            k = re.sub(
+                r'^_feature_blocks.res\d.block(\d)-(\d+)',
+                lambda x: f's{int(x.group(1))}.b{int(x.group(2)) + 1}', k)
+            k = re.sub(r's(\d)\.b(\d+)\.bn', r's\1.b\2.downsample.bn', k)
+            k = k.replace('proj', 'downsample.conv')
+            k = k.replace('f.a.0', 'conv1.conv')
+            k = k.replace('f.a.1', 'conv1.bn')
+            k = k.replace('f.b.0', 'conv2.conv')
+            k = k.replace('f.b.1', 'conv2.bn')
+            k = k.replace('f.c', 'conv3.conv')
+            k = k.replace('f.final_bn', 'conv3.bn')
+            k = k.replace('f.se.excitation.0', 'se.fc1')
+            k = k.replace('f.se.excitation.2', 'se.fc2')
+            out[k] = v
+        for k, v in state_dict['heads'].items():
+            if 'projection_head' in k or 'prototypes' in k:
+                continue
+            k = k.replace('0.clf.0', 'head.fc')
+            out[k] = v
+        return out
+
     if 'model' in state_dict:
         # For DeiT trained regnety_160 pretraiend model
         state_dict = state_dict['model']
@@ -666,6 +702,24 @@ def regnety_320(pretrained=False, **kwargs):
     return _create_regnet('regnety_320', pretrained, **kwargs)
 
 
+@register_model
+def regnety_640(pretrained=False, **kwargs):
+    """RegNetY-64GF"""
+    return _create_regnet('regnety_640', pretrained, **kwargs)
+
+
+@register_model
+def regnety_1280(pretrained=False, **kwargs):
+    """RegNetY-128GF"""
+    return _create_regnet('regnety_1280', pretrained, **kwargs)
+
+
+@register_model
+def regnety_2560(pretrained=False, **kwargs):
+    """RegNetY-256GF"""
+    return _create_regnet('regnety_2560', pretrained, **kwargs)
+
+
 @register_model
 def regnety_040s_gn(pretrained=False, **kwargs):
     """RegNetY-4.0GF w/ GroupNorm """

timm/models/resnetv2.py

@@ -57,52 +57,52 @@ default_cfgs = {
     # pretrained on imagenet21k, finetuned on imagenet1k
     'resnetv2_50x1_bitm': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R50x1-ILSVRC2012.npz',
-        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0),
+        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0, custom_load=True),
     'resnetv2_50x3_bitm': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R50x3-ILSVRC2012.npz',
-        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0),
+        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0, custom_load=True),
     'resnetv2_101x1_bitm': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R101x1-ILSVRC2012.npz',
-        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0),
+        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0, custom_load=True),
     'resnetv2_101x3_bitm': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R101x3-ILSVRC2012.npz',
-        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0),
+        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0, custom_load=True),
     'resnetv2_152x2_bitm': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R152x2-ILSVRC2012.npz',
-        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0),
+        input_size=(3, 448, 448), pool_size=(14, 14), crop_pct=1.0, custom_load=True),
     'resnetv2_152x4_bitm': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R152x4-ILSVRC2012.npz',
-        input_size=(3, 480, 480), pool_size=(15, 15), crop_pct=1.0),  # only one at 480x480?
+        input_size=(3, 480, 480), pool_size=(15, 15), crop_pct=1.0, custom_load=True),  # only one at 480x480?
 
     # trained on imagenet-21k
     'resnetv2_50x1_bitm_in21k': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R50x1.npz',
-        num_classes=21843),
+        num_classes=21843, custom_load=True),
     'resnetv2_50x3_bitm_in21k': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R50x3.npz',
-        num_classes=21843),
+        num_classes=21843, custom_load=True),
     'resnetv2_101x1_bitm_in21k': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R101x1.npz',
-        num_classes=21843),
+        num_classes=21843, custom_load=True),
     'resnetv2_101x3_bitm_in21k': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R101x3.npz',
-        num_classes=21843),
+        num_classes=21843, custom_load=True),
     'resnetv2_152x2_bitm_in21k': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R152x2.npz',
-        num_classes=21843),
+        num_classes=21843, custom_load=True),
     'resnetv2_152x4_bitm_in21k': _cfg(
         url='https://storage.googleapis.com/bit_models/BiT-M-R152x4.npz',
-        num_classes=21843),
+        num_classes=21843, custom_load=True),
 
     'resnetv2_50x1_bit_distilled': _cfg(
         url='https://storage.googleapis.com/bit_models/distill/R50x1_224.npz',
-        interpolation='bicubic'),
+        interpolation='bicubic', custom_load=True),
     'resnetv2_152x2_bit_teacher': _cfg(
         url='https://storage.googleapis.com/bit_models/distill/R152x2_T_224.npz',
-        interpolation='bicubic'),
+        interpolation='bicubic', custom_load=True),
     'resnetv2_152x2_bit_teacher_384': _cfg(
         url='https://storage.googleapis.com/bit_models/distill/R152x2_T_384.npz',
-        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, interpolation='bicubic'),
+        input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, interpolation='bicubic', custom_load=True),
 
     'resnetv2_50': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/resnetv2_50_a1h-000cdf49.pth',
@@ -507,8 +507,8 @@ def _create_resnetv2(variant, pretrained=False, **kwargs):
     return build_model_with_cfg(
         ResNetV2, variant, pretrained,
         feature_cfg=feature_cfg,
-        pretrained_custom_load='_bit' in variant,
-        **kwargs)
+        **kwargs,
+    )
 
 
 def _create_resnetv2_bit(variant, pretrained=False, **kwargs):

timm/models/vision_transformer.py

@@ -32,173 +32,14 @@ import torch.utils.checkpoint
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD,\
     OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
-from .helpers import build_model_with_cfg, resolve_pretrained_cfg, named_apply, adapt_input_conv, checkpoint_seq
+from .helpers import build_model_with_cfg, named_apply, adapt_input_conv, checkpoint_seq
 from .layers import PatchEmbed, Mlp, DropPath, trunc_normal_, lecun_normal_
+from .pretrained import generate_default_cfgs
 from .registry import register_model
 
 _logger = logging.getLogger(__name__)
 
 
-def _cfg(url='', **kwargs):
-    return {
-        'url': url,
-        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
-        'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,
-        'mean': IMAGENET_INCEPTION_MEAN, 'std': IMAGENET_INCEPTION_STD,
-        'first_conv': 'patch_embed.proj', 'classifier': 'head',
-        **kwargs
-    }
-
-
-default_cfgs = {
-    # patch models (weights from official Google JAX impl)
-    'vit_tiny_patch16_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz'),
-    'vit_tiny_patch16_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_small_patch32_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz'),
-    'vit_small_patch32_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_small_patch16_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz'),
-    'vit_small_patch16_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_base_patch32_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'B_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz'),
-    'vit_base_patch32_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_base_patch16_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npz'),
-    'vit_base_patch16_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_base_patch8_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npz'),
-    'vit_large_patch32_224': _cfg(
-        url='',  # no official model weights for this combo, only for in21k
-        ),
-    'vit_large_patch32_384': _cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p32_384-9b920ba8.pth',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_large_patch16_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npz'),
-    'vit_large_patch16_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-
-    'vit_large_patch14_224': _cfg(url=''),
-    'vit_huge_patch14_224': _cfg(url=''),
-    'vit_giant_patch14_224': _cfg(url=''),
-    'vit_gigantic_patch14_224': _cfg(url=''),
-
-
-    # patch models, imagenet21k (weights from official Google JAX impl)
-    'vit_tiny_patch16_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz',
-        num_classes=21843),
-    'vit_small_patch32_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npz',
-        num_classes=21843),
-    'vit_small_patch16_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npz',
-        num_classes=21843),
-    'vit_base_patch32_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.03-do_0.0-sd_0.0.npz',
-        num_classes=21843),
-    'vit_base_patch16_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz',
-        num_classes=21843),
-    'vit_base_patch8_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz',
-        num_classes=21843),
-    'vit_large_patch32_224_in21k': _cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth',
-        num_classes=21843),
-    'vit_large_patch16_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1.npz',
-        num_classes=21843),
-    'vit_huge_patch14_224_in21k': _cfg(
-        url='https://storage.googleapis.com/vit_models/imagenet21k/ViT-H_14.npz',
-        hf_hub_id='timm/vit_huge_patch14_224_in21k',
-        num_classes=21843),
-
-    # SAM trained models (https://arxiv.org/abs/2106.01548)
-    'vit_base_patch32_224_sam': _cfg(
-        url='https://storage.googleapis.com/vit_models/sam/ViT-B_32.npz'),
-    'vit_base_patch16_224_sam': _cfg(
-        url='https://storage.googleapis.com/vit_models/sam/ViT-B_16.npz'),
-
-    # DINO pretrained - https://arxiv.org/abs/2104.14294 (no classifier head, for fine-tune only)
-    'vit_small_patch16_224_dino': _cfg(
-        url='https://dl.fbaipublicfiles.com/dino/dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
-    'vit_small_patch8_224_dino': _cfg(
-        url='https://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
-    'vit_base_patch16_224_dino': _cfg(
-        url='https://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
-    'vit_base_patch8_224_dino': _cfg(
-        url='https://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth',
-        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
-
-
-    # ViT ImageNet-21K-P pretraining by MILL
-    'vit_base_patch16_224_miil_in21k': _cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/vit_base_patch16_224_in21k_miil-887286df.pth',
-        mean=(0., 0., 0.), std=(1., 1., 1.), crop_pct=0.875, interpolation='bilinear', num_classes=11221),
-    'vit_base_patch16_224_miil': _cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/vit_base_patch16_224_1k_miil_84_4-2deb18e3.pth',
-        mean=(0., 0., 0.), std=(1., 1., 1.), crop_pct=0.875, interpolation='bilinear'),
-
-    'vit_base_patch16_rpn_224': _cfg(
-        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/vit_base_patch16_rpn_224-sw-3b07e89d.pth'),
-
-    # experimental (may be removed)
-    'vit_base_patch32_plus_256': _cfg(url='', input_size=(3, 256, 256), crop_pct=0.95),
-    'vit_base_patch16_plus_240': _cfg(url='', input_size=(3, 240, 240), crop_pct=0.95),
-    'vit_small_patch16_36x1_224': _cfg(url=''),
-    'vit_small_patch16_18x2_224': _cfg(url=''),
-    'vit_base_patch16_18x2_224': _cfg(url=''),
-
-    'vit_base_patch32_224_clip_laion2b': _cfg(
-        hf_hub_id='laion/CLIP-ViT-B-32-laion2B-s34B-b79K',
-        hf_hub_filename='open_clip_pytorch_model.bin',
-        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=512),
-    'vit_large_patch14_224_clip_laion2b': _cfg(
-        hf_hub_id='laion/CLIP-ViT-L-14-laion2B-s32B-b82K',
-        hf_hub_filename='open_clip_pytorch_model.bin',
-        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD, num_classes=768),
-    'vit_huge_patch14_224_clip_laion2b': _cfg(
-        hf_hub_id='laion/CLIP-ViT-H-14-laion2B-s32B-b79K',
-        hf_hub_filename='open_clip_pytorch_model.bin',
-        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=1024),
-    'vit_giant_patch14_224_clip_laion2b': _cfg(
-        hf_hub_id='laion/CLIP-ViT-g-14-laion2B-s12B-b42K',
-        hf_hub_filename='open_clip_pytorch_model.bin',
-        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=1024),
-
-}
-
-
 class Attention(nn.Module):
     def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
         super().__init__()
@@ -651,7 +492,8 @@ def _load_weights(model: VisionTransformer, checkpoint_path: str, prefix: str =
             model.patch_embed.proj.weight.shape[1], _n2p(w[f'{prefix}embedding/kernel']))
     model.patch_embed.proj.weight.copy_(embed_conv_w)
     model.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))
-    model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
+    if model.cls_token is not None:
+        model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
     pos_embed_w = _n2p(w[f'{prefix}Transformer/posembed_input/pos_embedding'], t=False)
     if pos_embed_w.shape != model.pos_embed.shape:
         pos_embed_w = resize_pos_embed(  # resize pos embedding when different size from pretrained weights
@@ -778,18 +620,331 @@ def checkpoint_filter_fn(state_dict, model, adapt_layer_scale=False):
     return out_dict
 
 
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,
+        'mean': IMAGENET_INCEPTION_MEAN, 'std': IMAGENET_INCEPTION_STD,
+        'first_conv': 'patch_embed.proj', 'classifier': 'head',
+        **kwargs
+    }
+
+
+default_cfgs = generate_default_cfgs({
+
+    # How to train your ViT (augreg) weights, pretrained on 21k FT on in1k
+    'vit_tiny_patch16_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+        custom_load=True),
+    'vit_tiny_patch16_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+    'vit_small_patch32_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+        custom_load=True),
+    'vit_small_patch32_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+    'vit_small_patch16_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+        custom_load=True),
+    'vit_small_patch16_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+    'vit_base_patch32_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+        custom_load=True),
+    'vit_base_patch32_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+    'vit_base_patch16_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npz',
+        custom_load=True),
+    'vit_base_patch16_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+    'vit_base_patch8_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npz',
+        custom_load=True),
+    'vit_large_patch16_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npz',
+        custom_load=True),
+    'vit_large_patch16_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+
+    # re-finetuned augreg 21k FT on in1k weights
+    'vit_base_patch16_224.augreg2_in21k_ft_in1k': _cfg(
+        file='b16_augreg-a-8.pth'),
+    'vit_base_patch16_384.augreg2_in21k_ft_in1k': _cfg(
+        url=''),
+    'vit_base_patch8_224.augreg2_in21k_ft_in1k': _cfg(
+        url=''),
+
+    # patch models (weights from official Google JAX impl) pretrained on in21k FT on in1k
+    'vit_base_patch16_224.orig_in21k_ft_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth'),
+    'vit_base_patch16_384.orig_in21k_ft_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth'),
+    'vit_large_patch32_384.orig_in21k_ft_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p32_384-9b920ba8.pth',
+        input_size=(3, 384, 384), crop_pct=1.0),
+
+    # How to train your ViT (augreg) weights trained on in1k
+    'vit_base_patch16_224.augreg_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_16-i1k-300ep-lr_0.001-aug_strong2-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npz',
+        custom_load=True),
+    'vit_base_patch16_384.augreg_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_16-i1k-300ep-lr_0.001-aug_strong2-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz',
+        custom_load=True, input_size=(3, 384, 384), crop_pct=1.0),
+
+    'vit_large_patch14_224.untrained': _cfg(url=''),
+    'vit_huge_patch14_224.untrained': _cfg(url=''),
+    'vit_giant_patch14_224.untrained': _cfg(url=''),
+    'vit_gigantic_patch14_224.untrained': _cfg(url=''),
+
+
+    # patch models, imagenet21k (weights from official Google JAX impl)
+    'vit_large_patch32_224.v1_in21k': _cfg(
+            url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth',
+            num_classes=21843),
+    'vit_huge_patch14_224.v1_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/imagenet21k/ViT-H_14.npz',
+        hf_hub_id='timm/vit_huge_patch14_224_in21k',
+        custom_load=True, num_classes=21843),
+
+    # How to train your ViT (augreg) weights, pretrained on in21k
+    'vit_tiny_patch16_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz',
+        custom_load=True, num_classes=21843),
+    'vit_small_patch32_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npz',
+        custom_load=True, num_classes=21843),
+    'vit_small_patch16_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npz',
+        custom_load=True, num_classes=21843),
+    'vit_base_patch32_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.03-do_0.0-sd_0.0.npz',
+        custom_load=True, num_classes=21843),
+    'vit_base_patch16_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz',
+        custom_load=True, num_classes=21843),
+    'vit_base_patch8_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz',
+        custom_load=True, num_classes=21843),
+    'vit_large_patch16_224.augreg_in21k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1.npz',
+        custom_load=True, num_classes=21843),
+
+    # SAM trained models (https://arxiv.org/abs/2106.01548)
+    'vit_base_patch32_224.sam': _cfg(
+        url='https://storage.googleapis.com/vit_models/sam/ViT-B_32.npz', custom_load=True),
+    'vit_base_patch16_224.sam': _cfg(
+        url='https://storage.googleapis.com/vit_models/sam/ViT-B_16.npz', custom_load=True),
+
+    # DINO pretrained - https://arxiv.org/abs/2104.14294 (no classifier head, for fine-tune only)
+    'vit_small_patch16_224.dino': _cfg(
+        url='https://dl.fbaipublicfiles.com/dino/dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
+    'vit_small_patch8_224.dino': _cfg(
+        url='https://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
+    'vit_base_patch16_224.dino': _cfg(
+        url='https://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
+    'vit_base_patch8_224.dino': _cfg(
+        url='https://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth',
+        mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, num_classes=0),
+
+
+    # ViT ImageNet-21K-P pretraining by MILL
+    'vit_base_patch16_224_miil.in21k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/vit_base_patch16_224_in21k_miil-887286df.pth',
+        mean=(0., 0., 0.), std=(1., 1., 1.), crop_pct=0.875, interpolation='bilinear', num_classes=11221),
+    'vit_base_patch16_224_miil.in21k_ft_in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/vit_base_patch16_224_1k_miil_84_4-2deb18e3.pth',
+        mean=(0., 0., 0.), std=(1., 1., 1.), crop_pct=0.875, interpolation='bilinear'),
+
+    # custom timm variants
+    'vit_base_patch16_rpn_224.in1k': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/vit_base_patch16_rpn_224-sw-3b07e89d.pth'),
+    'vit_medium_patch16_gap_240.in12k': _cfg(
+        hf_hub_id='timm/vit_medium_patch16_gap_240.in12k',
+        input_size=(3, 240, 240), crop_pct=0.95, num_classes=11821),
+    'vit_medium_patch16_gap_256.in12k_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_medium_patch16_gap_256.in12k_ft_in1k',
+        input_size=(3, 256, 256), crop_pct=0.95),
+    'vit_medium_patch16_gap_384.in12k_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_medium_patch16_gap_384.in12k_ft_in1k',
+        input_size=(3, 384, 384), crop_pct=0.95, crop_mode='squash'),
+    'vit_base_patch16_gap_224': _cfg(),
+
+    # CLIP pretrained image tower and related fine-tuned weights
+    'vit_base_patch32_clip_224.laion2b': _cfg(
+        hf_hub_id='laion/CLIP-ViT-B-32-laion2B-s34B-b79K',
+        hf_hub_filename='open_clip_pytorch_model.bin',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=512),
+    'vit_base_patch16_clip_224.laion2b': _cfg(
+        #hf_hub_id='laion/CLIP-ViT-B-16-laion2B-s34B-b88K',
+        hf_hub_filename='open_clip_pytorch_model.bin',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, num_classes=512),
+    'vit_large_patch14_clip_224.laion2b': _cfg(
+        hf_hub_id='laion/CLIP-ViT-L-14-laion2B-s32B-b82K',
+        hf_hub_filename='open_clip_pytorch_model.bin',
+        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD, crop_pct=1.0, num_classes=768),
+    'vit_huge_patch14_clip_224.laion2b': _cfg(
+        hf_hub_id='laion/CLIP-ViT-H-14-laion2B-s32B-b79K',
+        hf_hub_filename='open_clip_pytorch_model.bin',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, num_classes=1024),
+    'vit_giant_patch14_clip_224.laion2b': _cfg(
+        hf_hub_id='laion/CLIP-ViT-g-14-laion2B-s12B-b42K',
+        hf_hub_filename='open_clip_pytorch_model.bin',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, num_classes=1024),
+
+    'vit_base_patch32_clip_224.laion2b_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch32_clip_224.laion2b_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD),
+    'vit_base_patch16_clip_224.laion2b_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_224.laion2b_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0),
+    'vit_base_patch16_clip_384.laion2b_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_384.laion2b_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=1.0, input_size=(3, 384, 384), crop_mode='squash'),
+    'vit_large_patch14_clip_224.laion2b_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_224.laion2b_ft_in1k',
+        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD, crop_pct=1.0),
+    'vit_large_patch14_clip_336.laion2b_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_336.laion2b_ft_in1k',
+        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD,
+        crop_pct=1.0, input_size=(3, 336, 336), crop_mode='squash'),
+    'vit_huge_patch14_clip_224.laion2b_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_huge_patch14_clip_224.laion2b_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0),
+    'vit_huge_patch14_clip_336.laion2b_ft_in1k': _cfg(
+        hf_hub_id='',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=1.0, input_size=(3, 336, 336), crop_mode='squash'),
+
+    'vit_base_patch32_clip_224.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch32_clip_224.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD),
+    'vit_base_patch32_clip_384.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch32_clip_384.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, input_size=(3, 384, 384)),
+    'vit_base_patch32_clip_448.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch32_clip_448.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, input_size=(3, 448, 448)),
+    'vit_base_patch16_clip_224.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_224.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=0.95),
+    'vit_base_patch16_clip_384.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_384.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=1.0, input_size=(3, 384, 384), crop_mode='squash'),
+    'vit_large_patch14_clip_224.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_224.laion2b_ft_in12k_in1k',
+        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD, crop_pct=1.0),
+    'vit_large_patch14_clip_336.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_336.laion2b_ft_in12k_in1k',
+        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD,
+        crop_pct=1.0, input_size=(3, 336, 336), crop_mode='squash'),
+    'vit_huge_patch14_clip_224.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_huge_patch14_clip_224.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0),
+    'vit_huge_patch14_clip_336.laion2b_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_huge_patch14_clip_336.laion2b_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=1.0, input_size=(3, 336, 336), crop_mode='squash'),
+
+    'vit_base_patch32_clip_224.laion2b_ft_in12k': _cfg(
+        #hf_hub_id='timm/vit_base_patch32_clip_224.laion2b_ft_in12k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=11821),
+    'vit_base_patch16_clip_224.laion2b_ft_in12k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_224.laion2b_ft_in12k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=11821),
+    'vit_large_patch14_clip_224.laion2b_ft_in12k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_224.laion2b_ft_in12k',
+        mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD, crop_pct=1.0, num_classes=11821),
+    'vit_huge_patch14_clip_224.laion2b_ft_in12k': _cfg(
+        hf_hub_id='timm/vit_huge_patch14_clip_224.laion2b_ft_in12k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, num_classes=11821),
+
+    'vit_base_patch32_clip_224.openai': _cfg(
+        hf_hub_id='timm/clip_vit_base_patch32_224.openai',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=512),
+    'vit_base_patch16_clip_224.openai': _cfg(
+        hf_hub_id='timm/clip_vit_base_patch16_224.openai',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=512),
+    'vit_large_patch14_clip_224.openai': _cfg(
+        hf_hub_id='timm/clip_vit_large_patch14_224.openai',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, num_classes=768),
+
+    'vit_base_patch32_clip_224.openai_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch32_clip_224.openai_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD),
+    'vit_base_patch16_clip_224.openai_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_224.openai_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD),
+    'vit_base_patch16_clip_384.openai_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_384.openai_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=1.0, input_size=(3, 384, 384), crop_mode='squash'),
+    'vit_large_patch14_clip_224.openai_ft_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_224.openai_ft_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0),
+
+    'vit_base_patch32_clip_224.openai_ft_in12k_in1k': _cfg(
+        #hf_hub_id='timm/vit_base_patch32_clip_224.openai_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD),
+    'vit_base_patch32_clip_384.openai_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch32_clip_384.openai_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=0.95, input_size=(3, 384, 384), crop_mode='squash'),
+    'vit_base_patch16_clip_224.openai_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_224.openai_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=0.95),
+    'vit_base_patch16_clip_384.openai_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_384.openai_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=0.95, input_size=(3, 384, 384), crop_mode='squash'),
+    'vit_large_patch14_clip_224.openai_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_224.openai_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0),
+    'vit_large_patch14_clip_336.openai_ft_in12k_in1k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_336.openai_ft_in12k_in1k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
+        crop_pct=1.0, input_size=(3, 336, 336), crop_mode='squash'),
+
+    'vit_base_patch32_clip_224.openai_ft_in12k': _cfg(
+        #hf_hub_id='timm/vit_base_patch32_clip_224.openai_ft_in12k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=11821),
+    'vit_base_patch16_clip_224.openai_ft_in12k': _cfg(
+        hf_hub_id='timm/vit_base_patch16_clip_224.openai_ft_in12k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, num_classes=11821),
+    'vit_large_patch14_clip_224.openai_ft_in12k': _cfg(
+        hf_hub_id='timm/vit_large_patch14_clip_224.openai_ft_in12k',
+        mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD, crop_pct=1.0, num_classes=11821),
+
+    # experimental (may be removed)
+    'vit_base_patch32_plus_256': _cfg(url='', input_size=(3, 256, 256), crop_pct=0.95),
+    'vit_base_patch16_plus_240': _cfg(url='', input_size=(3, 240, 240), crop_pct=0.95),
+    'vit_small_patch16_36x1_224': _cfg(url=''),
+    'vit_small_patch16_18x2_224': _cfg(url=''),
+    'vit_base_patch16_18x2_224': _cfg(url=''),
+})
+
+
 def _create_vision_transformer(variant, pretrained=False, **kwargs):
     if kwargs.get('features_only', None):
         raise RuntimeError('features_only not implemented for Vision Transformer models.')
 
-    pretrained_cfg = resolve_pretrained_cfg(variant, pretrained_cfg=kwargs.pop('pretrained_cfg', None))
-    model = build_model_with_cfg(
+    return build_model_with_cfg(
         VisionTransformer, variant, pretrained,
-        pretrained_cfg=pretrained_cfg,
         pretrained_filter_fn=checkpoint_filter_fn,
-        pretrained_custom_load='npz' in pretrained_cfg['url'],
-        **kwargs)
-    return model
+        **kwargs,
+    )
 
 
 @register_model
@@ -831,7 +986,6 @@ def vit_small_patch32_384(pretrained=False, **kwargs):
 @register_model
 def vit_small_patch16_224(pretrained=False, **kwargs):
     """ ViT-Small (ViT-S/16)
-    NOTE I've replaced my previous 'small' model definition and weights with the small variant from the DeiT paper
     """
     model_kwargs = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, **kwargs)
     model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **model_kwargs)
@@ -841,13 +995,21 @@ def vit_small_patch16_224(pretrained=False, **kwargs):
 @register_model
 def vit_small_patch16_384(pretrained=False, **kwargs):
     """ ViT-Small (ViT-S/16)
-    NOTE I've replaced my previous 'small' model definition and weights with the small variant from the DeiT paper
     """
     model_kwargs = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, **kwargs)
     model = _create_vision_transformer('vit_small_patch16_384', pretrained=pretrained, **model_kwargs)
     return model
 
 
+@register_model
+def vit_small_patch8_224(pretrained=False, **kwargs):
+    """ ViT-Small (ViT-S/8)
+    """
+    model_kwargs = dict(patch_size=8, embed_dim=384, depth=12, num_heads=6, **kwargs)
+    model = _create_vision_transformer('vit_small_patch8_224', pretrained=pretrained, **model_kwargs)
+    return model
+
+
 @register_model
 def vit_base_patch32_224(pretrained=False, **kwargs):
     """ ViT-Base (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
@@ -974,175 +1136,158 @@ def vit_gigantic_patch14_224(pretrained=False, **kwargs):
 
 
 @register_model
-def vit_tiny_patch16_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Tiny (Vit-Ti/16).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
-    """
-    model_kwargs = dict(patch_size=16, embed_dim=192, depth=12, num_heads=3, **kwargs)
-    model = _create_vision_transformer('vit_tiny_patch16_224_in21k', pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_small_patch32_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Small (ViT-S/16)
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
-    """
-    model_kwargs = dict(patch_size=32, embed_dim=384, depth=12, num_heads=6, **kwargs)
-    model = _create_vision_transformer('vit_small_patch32_224_in21k', pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_small_patch16_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Small (ViT-S/16)
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
+def vit_base_patch16_224_miil(pretrained=False, **kwargs):
+    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
+    Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
     """
-    model_kwargs = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, **kwargs)
-    model = _create_vision_transformer('vit_small_patch16_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, qkv_bias=False, **kwargs)
+    model = _create_vision_transformer('vit_base_patch16_224_miil', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch32_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
+def vit_medium_patch16_gap_240(pretrained=False, **kwargs):
+    """ ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 240x240
     """
-    model_kwargs = dict(patch_size=32, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch32_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=16, embed_dim=512, depth=12, num_heads=8, class_token=False,
+        global_pool=kwargs.get('global_pool', 'avg'), qkv_bias=False, init_values=1e-6, fc_norm=False, **kwargs)
+    model = _create_vision_transformer('vit_medium_patch16_gap_240', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch16_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
+def vit_medium_patch16_gap_256(pretrained=False, **kwargs):
+    """ ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 256x256
     """
-    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch16_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=16, embed_dim=512, depth=12, num_heads=8, class_token=False,
+        global_pool=kwargs.get('global_pool', 'avg'), qkv_bias=False, init_values=1e-6, fc_norm=False, **kwargs)
+    model = _create_vision_transformer('vit_medium_patch16_gap_256', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch8_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Base model (ViT-B/8) from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
+def vit_medium_patch16_gap_384(pretrained=False, **kwargs):
+    """ ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 384x384
     """
-    model_kwargs = dict(patch_size=8, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch8_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=16, embed_dim=512, depth=12, num_heads=8, class_token=False,
+        global_pool=kwargs.get('global_pool', 'avg'), qkv_bias=False, init_values=1e-6, fc_norm=False, **kwargs)
+    model = _create_vision_transformer('vit_medium_patch16_gap_384', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_large_patch32_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has a representation layer but the 21k classifier head is zero'd out in original weights
+def vit_base_patch16_gap_224(pretrained=False, **kwargs):
+    """ ViT-Base (ViT-B/16) w/o class token, w/ avg-pool @ 256x256
     """
-    model_kwargs = dict(patch_size=32, embed_dim=1024, depth=24, num_heads=16, **kwargs)
-    model = _create_vision_transformer('vit_large_patch32_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=16, class_token=False,
+        global_pool=kwargs.get('global_pool', 'avg'), fc_norm=False, **kwargs)
+    model = _create_vision_transformer('vit_base_patch16_gap_224', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_large_patch16_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has valid 21k classifier head and no representation (pre-logits) layer
+def vit_base_patch32_clip_224(pretrained=False, **kwargs):
+    """ ViT-B/32 CLIP image tower @ 224x224
     """
-    model_kwargs = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs)
-    model = _create_vision_transformer('vit_large_patch16_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=32, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_base_patch32_clip_224', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_huge_patch14_224_in21k(pretrained=False, **kwargs):
-    """ ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    NOTE: this model has a representation layer but the 21k classifier head is zero'd out in original weights
+def vit_base_patch32_clip_384(pretrained=False, **kwargs):
+    """ ViT-B/32 CLIP image tower @ 384x384
     """
-    model_kwargs = dict(patch_size=14, embed_dim=1280, depth=32, num_heads=16, **kwargs)
-    model = _create_vision_transformer('vit_huge_patch14_224_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=32, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_base_patch32_clip_384', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch16_224_sam(pretrained=False, **kwargs):
-    """ ViT-Base (ViT-B/16) w/ SAM pretrained weights. Paper: https://arxiv.org/abs/2106.01548
+def vit_base_patch32_clip_448(pretrained=False, **kwargs):
+    """ ViT-B/32 CLIP image tower @ 448x448
     """
-    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch16_224_sam', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=32, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_base_patch32_clip_448', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch32_224_sam(pretrained=False, **kwargs):
-    """ ViT-Base (ViT-B/32) w/ SAM pretrained weights. Paper: https://arxiv.org/abs/2106.01548
+def vit_base_patch16_clip_224(pretrained=False, **kwargs):
+    """ ViT-B/16 CLIP image tower
     """
-    model_kwargs = dict(patch_size=32, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch32_224_sam', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_base_patch16_clip_224', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_small_patch16_224_dino(pretrained=False, **kwargs):
-    """ ViT-Small (ViT-S/16) w/ DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294
+def vit_base_patch16_clip_384(pretrained=False, **kwargs):
+    """ ViT-B/16 CLIP image tower @ 384x384
     """
-    model_kwargs = dict(patch_size=16, embed_dim=384, depth=12, num_heads=6, **kwargs)
-    model = _create_vision_transformer('vit_small_patch16_224_dino', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_base_patch16_clip_384', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_small_patch8_224_dino(pretrained=False, **kwargs):
-    """ ViT-Small (ViT-S/8) w/ DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294
+def vit_large_patch14_clip_224(pretrained=False, **kwargs):
+    """ ViT-Large model (ViT-L/14) CLIP image tower
     """
-    model_kwargs = dict(patch_size=8, embed_dim=384, depth=12, num_heads=6, **kwargs)
-    model = _create_vision_transformer('vit_small_patch8_224_dino', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=14, embed_dim=1024, depth=24, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_large_patch14_clip_224', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch16_224_dino(pretrained=False, **kwargs):
-    """ ViT-Base (ViT-B/16) /w DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294
+def vit_large_patch14_clip_336(pretrained=False, **kwargs):
+    """ ViT-Large model (ViT-L/14) CLIP image tower @ 336x336
     """
-    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch16_224_dino', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=14, embed_dim=1024, depth=24, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_large_patch14_clip_336', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch8_224_dino(pretrained=False, **kwargs):
-    """ ViT-Base (ViT-B/8) w/ DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294
+def vit_huge_patch14_clip_224(pretrained=False, **kwargs):
+    """ ViT-Huge model (ViT-H/14) CLIP image tower.
     """
-    model_kwargs = dict(patch_size=8, embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer('vit_base_patch8_224_dino', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_huge_patch14_clip_224', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch16_224_miil_in21k(pretrained=False, **kwargs):
-    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
-    Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
+def vit_huge_patch14_clip_336(pretrained=False, **kwargs):
+    """ ViT-Huge model (ViT-H/14) CLIP image tower @ 336x336
     """
-    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, qkv_bias=False, **kwargs)
-    model = _create_vision_transformer('vit_base_patch16_224_miil_in21k', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_huge_patch14_clip_336', pretrained=pretrained, **model_kwargs)
     return model
 
 
 @register_model
-def vit_base_patch16_224_miil(pretrained=False, **kwargs):
-    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
-    Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
+def vit_giant_patch14_clip_224(pretrained=False, **kwargs):
+    """ ViT-Giant (little-g) model (ViT-g/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
+    Pretrained weights from CLIP image tower.
     """
-    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, qkv_bias=False, **kwargs)
-    model = _create_vision_transformer('vit_base_patch16_224_miil', pretrained=pretrained, **model_kwargs)
+    model_kwargs = dict(
+        patch_size=14, embed_dim=1408, mlp_ratio=48/11, depth=40, num_heads=16,
+        pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
+    model = _create_vision_transformer('vit_giant_patch14_clip_224', pretrained=pretrained, **model_kwargs)
     return model
 
 
@@ -1209,48 +1354,3 @@ def vit_base_patch16_18x2_224(pretrained=False, **kwargs):
         patch_size=16, embed_dim=768, depth=18, num_heads=12, init_values=1e-5, block_fn=ParallelBlock, **kwargs)
     model = _create_vision_transformer('vit_base_patch16_18x2_224', pretrained=pretrained, **model_kwargs)
     return model
-
-
-@register_model
-def vit_base_patch32_224_clip_laion2b(pretrained=False, **kwargs):
-    """ ViT-B/32
-    Pretrained weights from CLIP image tower trained on LAION-2B image-text pairs.
-    """
-    model_kwargs = dict(
-        patch_size=32, embed_dim=768, depth=12, num_heads=12, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
-    model = _create_vision_transformer('vit_base_patch32_224_clip_laion2b', pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_large_patch14_224_clip_laion2b(pretrained=False, **kwargs):
-    """ ViT-Large model (ViT-L/14)
-    Pretrained weights from CLIP image tower trained on LAION-2B image-text pairs.
-    """
-    model_kwargs = dict(
-        patch_size=14, embed_dim=1024, depth=24, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
-    model = _create_vision_transformer('vit_large_patch14_224_clip_laion2b', pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_huge_patch14_224_clip_laion2b(pretrained=False, **kwargs):
-    """ ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).
-    Pretrained weights from CLIP image tower trained on LAION-2B image-text pairs.
-    """
-    model_kwargs = dict(
-        patch_size=14, embed_dim=1280, depth=32, num_heads=16, pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
-    model = _create_vision_transformer('vit_huge_patch14_224_clip_laion2b', pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_giant_patch14_224_clip_laion2b(pretrained=False, **kwargs):
-    """ ViT-Giant (little-g) model (ViT-g/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
-    Pretrained weights from CLIP image tower trained on LAION-2B image-text pairs.
-    """
-    model_kwargs = dict(
-        patch_size=14, embed_dim=1408, mlp_ratio=48/11, depth=40, num_heads=16,
-        pre_norm=True, norm_layer=nn.LayerNorm, **kwargs)
-    model = _create_vision_transformer('vit_giant_patch14_224_clip_laion2b', pretrained=pretrained, **model_kwargs)
-    return model

timm/models/vision_transformer_hybrid.py

@@ -21,6 +21,7 @@ import torch.nn as nn
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from .layers import StdConv2dSame, StdConv2d, to_2tuple
+from .pretrained import generate_default_cfgs
 from .resnet import resnet26d, resnet50d
 from .resnetv2 import ResNetV2, create_resnetv2_stem
 from .registry import register_model
@@ -38,52 +39,48 @@ def _cfg(url='', **kwargs):
     }
 
 
-default_cfgs = {
+default_cfgs = generate_default_cfgs({
     # hybrid in-1k models (weights from official JAX impl where they exist)
-    'vit_tiny_r_s16_p8_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+    'vit_tiny_r_s16_p8_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+        custom_load=True,
         first_conv='patch_embed.backbone.conv'),
-    'vit_tiny_r_s16_p8_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
-        first_conv='patch_embed.backbone.conv', input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_small_r26_s32_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'R26_S_32-i21k-300ep-lr_0.001-aug_light0-wd_0.03-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+    'vit_tiny_r_s16_p8_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
+        first_conv='patch_embed.backbone.conv', input_size=(3, 384, 384), crop_pct=1.0, custom_load=True),
+    'vit_small_r26_s32_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_light0-wd_0.03-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.03-res_224.npz',
+        custom_load=True,
     ),
-    'vit_small_r26_s32_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'R26_S_32-i21k-300ep-lr_0.001-aug_medium2-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_base_r26_s32_224': _cfg(),
-    'vit_base_r50_s16_224': _cfg(),
-    'vit_base_r50_s16_384': _cfg(
+    'vit_small_r26_s32_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_medium2-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz',
+        input_size=(3, 384, 384), crop_pct=1.0, custom_load=True),
+    'vit_base_r26_s32_224.untrained': _cfg(),
+    'vit_base_r50_s16_384.v1_in21k_ft_in1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_resnet50_384-9fd3c705.pth',
         input_size=(3, 384, 384), crop_pct=1.0),
-    'vit_large_r50_s32_224': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npz'
+    'vit_large_r50_s32_224.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npz',
+        custom_load=True,
     ),
-    'vit_large_r50_s32_384': _cfg(
-        url='https://storage.googleapis.com/vit_models/augreg/'
-            'R50_L_32-i21k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz',
-        input_size=(3, 384, 384), crop_pct=1.0
+    'vit_large_r50_s32_384.augreg_in21k_ft_in1k': _cfg(
+        url='https://storage.googleapis.com/vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz',
+        input_size=(3, 384, 384), crop_pct=1.0, custom_load=True,
     ),
 
     # hybrid in-21k models (weights from official Google JAX impl where they exist)
-    'vit_tiny_r_s16_p8_224_in21k': _cfg(
+    'vit_tiny_r_s16_p8_224.augreg_in21k': _cfg(
         url='https://storage.googleapis.com/vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz',
-        num_classes=21843, crop_pct=0.9, first_conv='patch_embed.backbone.conv'),
-    'vit_small_r26_s32_224_in21k': _cfg(
+        num_classes=21843, crop_pct=0.9, first_conv='patch_embed.backbone.conv', custom_load=True),
+    'vit_small_r26_s32_224.augreg_in21k': _cfg(
         url='https://storage.googleapis.com/vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_medium2-wd_0.03-do_0.0-sd_0.0.npz',
-        num_classes=21843, crop_pct=0.9),
-    'vit_base_r50_s16_224_in21k': _cfg(
+        num_classes=21843, crop_pct=0.9, custom_load=True),
+    'vit_base_r50_s16_224.v1_in21k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_resnet50_224_in21k-6f7c7740.pth',
         num_classes=21843, crop_pct=0.9),
-    'vit_large_r50_s32_224_in21k': _cfg(
+    'vit_large_r50_s32_224.augreg_in21k': _cfg(
         url='https://storage.googleapis.com/vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0.npz',
-        num_classes=21843, crop_pct=0.9),
+        num_classes=21843, crop_pct=0.9, custom_load=True),
 
     # hybrid models (using timm resnet backbones)
     'vit_small_resnet26d_224': _cfg(
@@ -94,7 +91,7 @@ default_cfgs = {
         mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, first_conv='patch_embed.backbone.conv1.0'),
     'vit_base_resnet50d_224': _cfg(
         mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD, first_conv='patch_embed.backbone.conv1.0'),
-}
+})
 
 
 class HybridEmbed(nn.Module):
@@ -248,12 +245,6 @@ def vit_base_r50_s16_384(pretrained=False, **kwargs):
     return model
 
 
-@register_model
-def vit_base_resnet50_384(pretrained=False, **kwargs):
-    # DEPRECATED this is forwarding to model def above for backwards compatibility
-    return vit_base_r50_s16_384(pretrained=pretrained, **kwargs)
-
-
 @register_model
 def vit_large_r50_s32_224(pretrained=False, **kwargs):
     """ R50+ViT-L/S32 hybrid.
@@ -276,57 +267,6 @@ def vit_large_r50_s32_384(pretrained=False, **kwargs):
     return model
 
 
-@register_model
-def vit_tiny_r_s16_p8_224_in21k(pretrained=False, **kwargs):
-    """ R+ViT-Ti/S16 w/ 8x8 patch hybrid.  ImageNet-21k.
-    """
-    backbone = _resnetv2(layers=(), **kwargs)
-    model_kwargs = dict(patch_size=8, embed_dim=192, depth=12, num_heads=3, **kwargs)
-    model = _create_vision_transformer_hybrid(
-        'vit_tiny_r_s16_p8_224_in21k', backbone=backbone, pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_small_r26_s32_224_in21k(pretrained=False, **kwargs):
-    """ R26+ViT-S/S32 hybrid. ImageNet-21k.
-    """
-    backbone = _resnetv2((2, 2, 2, 2), **kwargs)
-    model_kwargs = dict(embed_dim=384, depth=12, num_heads=6, **kwargs)
-    model = _create_vision_transformer_hybrid(
-        'vit_small_r26_s32_224_in21k', backbone=backbone, pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_base_r50_s16_224_in21k(pretrained=False, **kwargs):
-    """ R50+ViT-B/16 hybrid model from original paper (https://arxiv.org/abs/2010.11929).
-    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
-    """
-    backbone = _resnetv2(layers=(3, 4, 9), **kwargs)
-    model_kwargs = dict(embed_dim=768, depth=12, num_heads=12, **kwargs)
-    model = _create_vision_transformer_hybrid(
-        'vit_base_r50_s16_224_in21k', backbone=backbone, pretrained=pretrained, **model_kwargs)
-    return model
-
-
-@register_model
-def vit_base_resnet50_224_in21k(pretrained=False, **kwargs):
-    # DEPRECATED this is forwarding to model def above for backwards compatibility
-    return vit_base_r50_s16_224_in21k(pretrained=pretrained, **kwargs)
-
-
-@register_model
-def vit_large_r50_s32_224_in21k(pretrained=False, **kwargs):
-    """ R50+ViT-L/S32 hybrid. ImageNet-21k.
-    """
-    backbone = _resnetv2((3, 4, 6, 3), **kwargs)
-    model_kwargs = dict(embed_dim=1024, depth=24, num_heads=16, **kwargs)
-    model = _create_vision_transformer_hybrid(
-        'vit_large_r50_s32_224_in21k', backbone=backbone, pretrained=pretrained, **model_kwargs)
-    return model
-
-
 @register_model
 def vit_small_resnet26d_224(pretrained=False, **kwargs):
     """ Custom ViT small hybrid w/ ResNet26D stride 32. No pretrained weights.

timm/version.py

@@ -1 +1 @@
-__version__ = '0.7.0dev0'
+__version__ = '0.8.0dev0'

train.py

@@ -66,6 +66,8 @@ try:
 except ImportError as e:
     has_functorch = False
 
+has_compile = hasattr(torch, 'compile')
+
 
 _logger = logging.getLogger('train')
 
@@ -81,10 +83,12 @@ parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
 # Dataset parameters
 group = parser.add_argument_group('Dataset parameters')
 # Keep this argument outside of the dataset group because it is positional.
-parser.add_argument('data_dir', metavar='DIR',
-                    help='path to dataset')
-group.add_argument('--dataset', '-d', metavar='NAME', default='',
-                    help='dataset type (default: ImageFolder/ImageTar if empty)')
+parser.add_argument('data', nargs='?', metavar='DIR', const=None,
+                    help='path to dataset (positional is *deprecated*, use --data-dir)')
+parser.add_argument('--data-dir', metavar='DIR',
+                    help='path to dataset (root dir)')
+parser.add_argument('--dataset', metavar='NAME', default='',
+                    help='dataset type + name ("<type>/<name>") (default: ImageFolder or ImageTar if empty)')
 group.add_argument('--train-split', metavar='NAME', default='train',
                     help='dataset train split (default: train)')
 group.add_argument('--val-split', metavar='NAME', default='validation',
@@ -130,17 +134,20 @@ group.add_argument('-vb', '--validation-batch-size', type=int, default=None, met
                     help='Validation batch size override (default: None)')
 group.add_argument('--channels-last', action='store_true', default=False,
                     help='Use channels_last memory layout')
-scripting_group = group.add_mutually_exclusive_group()
-scripting_group.add_argument('--torchscript', dest='torchscript', action='store_true',
-                    help='torch.jit.script the full model')
-scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
-                    help="Enable AOT Autograd support. (It's recommended to use this option with `--fuser nvfuser` together)")
 group.add_argument('--fuser', default='', type=str,
                     help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
-group.add_argument('--fast-norm', default=False, action='store_true',
-                    help='enable experimental fast-norm')
 group.add_argument('--grad-checkpointing', action='store_true', default=False,
                     help='Enable gradient checkpointing through model blocks/stages')
+group.add_argument('--fast-norm', default=False, action='store_true',
+                    help='enable experimental fast-norm')
+
+scripting_group = group.add_mutually_exclusive_group()
+scripting_group.add_argument('--torchscript', dest='torchscript', action='store_true',
+                             help='torch.jit.script the full model')
+scripting_group.add_argument('--torchcompile', nargs='?', type=str, default=None, const='inductor',
+                             help="Enable compilation w/ specified backend (default: inductor).")
+scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
+                             help="Enable AOT Autograd support.")
 
 # Optimizer parameters
 group = parser.add_argument_group('Optimizer parameters')
@@ -365,6 +372,8 @@ def main():
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.benchmark = True
 
+    if args.data and not args.data_dir:
+        args.data_dir = args.data
     args.prefetcher = not args.no_prefetcher
     device = utils.init_distributed_device(args)
     if args.distributed:
@@ -473,12 +482,16 @@ def main():
         assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
         assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
         model = torch.jit.script(model)
-
-    if args.aot_autograd:
+    elif args.torchcompile:
+        # FIXME dynamo might need move below DDP wrapping? TBD
+        assert has_compile, 'A version of torch w/ torch.compile() is required for --compile, possibly a nightly.'
+        torch._dynamo.reset()
+        model = torch.compile(model, backend=args.torchcompile)
+    elif args.aot_autograd:
         assert has_functorch, "functorch is needed for --aot-autograd"
         model = memory_efficient_fusion(model)
 
-    if args.lr is None:
+    if not args.lr:
         global_batch_size = args.batch_size * args.world_size
         batch_ratio = global_batch_size / args.lr_base_size
         if not args.lr_base_scale:

validate.py

@@ -26,12 +26,11 @@ from timm.data import create_dataset, create_loader, resolve_data_config, RealLa
 from timm.utils import accuracy, AverageMeter, natural_key, setup_default_logging, set_jit_fuser,\
     decay_batch_step, check_batch_size_retry
 
-has_apex = False
 try:
     from apex import amp
     has_apex = True
 except ImportError:
-    pass
+    has_apex = False
 
 has_native_amp = False
 try:
@@ -46,14 +45,18 @@ try:
 except ImportError as e:
     has_functorch = False
 
+has_compile = hasattr(torch, 'compile')
+
 _logger = logging.getLogger('validate')
 
 
 parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
-parser.add_argument('data', metavar='DIR',
-                    help='path to dataset')
-parser.add_argument('--dataset', '-d', metavar='NAME', default='',
-                    help='dataset type (default: ImageFolder/ImageTar if empty)')
+parser.add_argument('data', nargs='?', metavar='DIR', const=None,
+                    help='path to dataset (*deprecated*, use --data-dir)')
+parser.add_argument('--data-dir', metavar='DIR',
+                    help='path to dataset (root dir)')
+parser.add_argument('--dataset', metavar='NAME', default='',
+                    help='dataset type + name ("<type>/<name>") (default: ImageFolder or ImageTar if empty)')
 parser.add_argument('--split', metavar='NAME', default='validation',
                     help='dataset split (default: validation)')
 parser.add_argument('--dataset-download', action='store_true', default=False,
@@ -72,6 +75,8 @@ parser.add_argument('--use-train-size', action='store_true', default=False,
                     help='force use of train input size, even when test size is specified in pretrained cfg')
 parser.add_argument('--crop-pct', default=None, type=float,
                     metavar='N', help='Input image center crop pct')
+parser.add_argument('--crop-mode', default=None, type=str,
+                    metavar='N', help='Input image crop mode (squash, border, center). Model default if None.')
 parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
                     help='Override mean pixel value of dataset')
 parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
@@ -112,17 +117,23 @@ parser.add_argument('--tf-preprocessing', action='store_true', default=False,
                     help='Use Tensorflow preprocessing pipeline (require CPU TF installed')
 parser.add_argument('--use-ema', dest='use_ema', action='store_true',
                     help='use ema version of weights if present')
-scripting_group = parser.add_mutually_exclusive_group()
-scripting_group.add_argument('--torchscript', dest='torchscript', action='store_true',
-                    help='torch.jit.script the full model')
-scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
-                    help="Enable AOT Autograd support. (It's recommended to use this option with `--fuser nvfuser` together)")
 parser.add_argument('--fuser', default='', type=str,
                     help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
 parser.add_argument('--fast-norm', default=False, action='store_true',
                     help='enable experimental fast-norm')
+
+scripting_group = parser.add_mutually_exclusive_group()
+scripting_group.add_argument('--torchscript', default=False, action='store_true',
+                             help='torch.jit.script the full model')
+scripting_group.add_argument('--torchcompile', nargs='?', type=str, default=None, const='inductor',
+                             help="Enable compilation w/ specified backend (default: inductor).")
+scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
+                             help="Enable AOT Autograd support.")
+
 parser.add_argument('--results-file', default='', type=str, metavar='FILENAME',
                     help='Output csv file for validation results (summary)')
+parser.add_argument('--results-format', default='csv', type=str,
+                    help='Format for results file one of (csv, json) (default: csv).')
 parser.add_argument('--real-labels', default='', type=str, metavar='FILENAME',
                     help='Real labels JSON file for imagenet evaluation')
 parser.add_argument('--valid-labels', default='', type=str, metavar='FILENAME',
@@ -196,21 +207,24 @@ def validate(args):
     if args.test_pool:
         model, test_time_pool = apply_test_time_pool(model, data_config)
 
+    model = model.to(device)
+    if args.channels_last:
+        model = model.to(memory_format=torch.channels_last)
+
     if args.torchscript:
-        torch.jit.optimized_execution(True)
+        assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
         model = torch.jit.script(model)
-
-    if args.aot_autograd:
+    elif args.torchcompile:
+        assert has_compile, 'A version of torch w/ torch.compile() is required for --compile, possibly a nightly.'
+        torch._dynamo.reset()
+        model = torch.compile(model, backend=args.torchcompile)
+    elif args.aot_autograd:
         assert has_functorch, "functorch is needed for --aot-autograd"
         model = memory_efficient_fusion(model)
 
-    model = model.to(device)
     if use_amp == 'apex':
         model = amp.initialize(model, opt_level='O1')
 
-    if args.channels_last:
-        model = model.to(memory_format=torch.channels_last)
-
     if args.num_gpu > 1:
         model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))
 
@@ -248,6 +262,7 @@ def validate(args):
         std=data_config['std'],
         num_workers=args.workers,
         crop_pct=crop_pct,
+        crop_mode=data_config['crop_mode'],
         pin_memory=args.pin_mem,
         device=device,
         tf_preprocessing=args.tf_preprocessing,
@@ -376,7 +391,7 @@ def main():
             model_cfgs = [(n, '') for n in model_names]
         elif not is_model(args.model):
             # model name doesn't exist, try as wildcard filter
-            model_names = list_models(args.model)
+            model_names = list_models(args.model, pretrained=True)
             model_cfgs = [(n, '') for n in model_names]
 
         if not model_cfgs and os.path.isfile(args.model):
@@ -385,7 +400,6 @@ def main():
             model_cfgs = [(n, None) for n in model_names if n]
 
     if len(model_cfgs):
-        results_file = args.results_file or './results-all.csv'
         _logger.info('Running bulk validation on these pretrained models: {}'.format(', '.join(model_names)))
         results = []
         try:
@@ -402,24 +416,34 @@ def main():
         except KeyboardInterrupt as e:
             pass
         results = sorted(results, key=lambda x: x['top1'], reverse=True)
-        if len(results):
-            write_results(results_file, results)
     else:
         if args.retry:
             results = _try_run(args, args.batch_size)
         else:
             results = validate(args)
+
+    if args.results_file:
+        write_results(args.results_file, results, format=args.results_format)
+
     # output results in JSON to stdout w/ delimiter for runner script
     print(f'--result\n{json.dumps(results, indent=4)}')
 
 
-def write_results(results_file, results):
+def write_results(results_file, results, format='csv'):
     with open(results_file, mode='w') as cf:
-        dw = csv.DictWriter(cf, fieldnames=results[0].keys())
-        dw.writeheader()
-        for r in results:
-            dw.writerow(r)
-        cf.flush()
+        if format == 'json':
+            json.dump(results, cf, indent=4)
+        else:
+            if not isinstance(results, (list, tuple)):
+                results = [results]
+            if not results:
+                return
+            dw = csv.DictWriter(cf, fieldnames=results[0].keys())
+            dw.writeheader()
+            for r in results:
+                dw.writerow(r)
+            cf.flush()
+
 
 
 if __name__ == '__main__':