Merge pull request #1327 from rwightman/edgenext_csp_and_more

EdgeNeXt, additional DarkNets, and more

Author: rwightman

benchmark.py

@@ -6,24 +6,23 @@
 Hacked together by Ross Wightman (https://github.com/rwightman)
 """
 import argparse
-import os
 import csv
 import json
-import time
 import logging
-import torch
-import torch.nn as nn
-import torch.nn.parallel
+import time
 from collections import OrderedDict
 from contextlib import suppress
 from functools import partial
 
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+
+from timm.data import resolve_data_config
 from timm.models import create_model, is_model, list_models
 from timm.optim import create_optimizer_v2
-from timm.data import resolve_data_config
 from timm.utils import setup_default_logging, set_jit_fuser
 
-
 has_apex = False
 try:
     from apex import amp
@@ -71,6 +70,8 @@
                     help="Benchmark mode. One of 'inference', 'train', 'both'. Defaults to 'both'")
 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',
                     help='Output csv file for validation results (summary)')
 parser.add_argument('--num-warm-iter', default=10, type=int,
@@ -169,10 +170,9 @@ def resolve_precision(precision: str):
 
 
 def profile_deepspeed(model, input_size=(3, 224, 224), batch_size=1, detailed=False):
-    macs, _ = get_model_profile(
+    _, macs, _ = get_model_profile(
         model=model,
-        input_res=(batch_size,) + input_size,  # input shape or input to the input_constructor
-        input_constructor=None,  # if specified, a constructor taking input_res is used as input to the model
+        input_shape=(batch_size,) + input_size,  # input shape/resolution
         print_profile=detailed,  # prints the model graph with the measured profile attached to each module
         detailed=detailed,  # print the detailed profile
         warm_up=10,  # the number of warm-ups before measuring the time of each module
@@ -197,8 +197,19 @@ def profile_fvcore(model, input_size=(3, 224, 224), batch_size=1, detailed=False
 
 class BenchmarkRunner:
     def __init__(
-            self, model_name, detail=False, device='cuda', torchscript=False, aot_autograd=False, precision='float32',
-            fuser='', num_warm_iter=10, num_bench_iter=50, use_train_size=False, **kwargs):
+            self,
+            model_name,
+            detail=False,
+            device='cuda',
+            torchscript=False,
+            aot_autograd=False,
+            precision='float32',
+            fuser='',
+            num_warm_iter=10,
+            num_bench_iter=50,
+            use_train_size=False,
+            **kwargs
+    ):
         self.model_name = model_name
         self.detail = detail
         self.device = device
@@ -225,11 +236,12 @@ def __init__(
         self.num_classes = self.model.num_classes
         self.param_count = count_params(self.model)
         _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
-        data_config = resolve_data_config(kwargs, model=self.model, use_test_size=not use_train_size)
         self.input_size = data_config['input_size']
         self.batch_size = kwargs.pop('batch_size', 256)
 
@@ -255,7 +267,13 @@ def _init_input(self):
 
 class InferenceBenchmarkRunner(BenchmarkRunner):
 
-    def __init__(self, model_name, device='cuda', torchscript=False, **kwargs):
+    def __init__(
+            self,
+            model_name,
+            device='cuda',
+            torchscript=False,
+            **kwargs
+    ):
         super().__init__(model_name=model_name, device=device, torchscript=torchscript, **kwargs)
         self.model.eval()
 
@@ -324,7 +342,13 @@ def _step():
 
 class TrainBenchmarkRunner(BenchmarkRunner):
 
-    def __init__(self, model_name, device='cuda', torchscript=False, **kwargs):
+    def __init__(
+            self,
+            model_name,
+            device='cuda',
+            torchscript=False,
+            **kwargs
+    ):
         super().__init__(model_name=model_name, device=device, torchscript=torchscript, **kwargs)
         self.model.train()
 
@@ -491,7 +515,7 @@ def decay_batch_exp(batch_size, factor=0.5, divisor=16):
     return max(0, int(out_batch_size))
 
 
-def _try_run(model_name, bench_fn, initial_batch_size, bench_kwargs):
+def _try_run(model_name, bench_fn, bench_kwargs, initial_batch_size, no_batch_size_retry=False):
     batch_size = initial_batch_size
     results = dict()
     error_str = 'Unknown'
@@ -506,8 +530,11 @@ def _try_run(model_name, bench_fn, initial_batch_size, bench_kwargs):
             if 'channels_last' in error_str:
                 _logger.error(f'{model_name} not supported in channels_last, skipping.')
                 break
-            _logger.warning(f'"{error_str}" while running benchmark. Reducing batch size to {batch_size} for retry.')
+            _logger.error(f'"{error_str}" while running benchmark.')
+            if no_batch_size_retry:
+                break
         batch_size = decay_batch_exp(batch_size)
+        _logger.warning(f'Reducing batch size to {batch_size} for retry.')
     results['error'] = error_str
     return results
 
@@ -549,7 +576,13 @@ def benchmark(args):
 
     model_results = OrderedDict(model=model)
     for prefix, bench_fn in zip(prefixes, bench_fns):
-        run_results = _try_run(model, bench_fn, initial_batch_size=batch_size, bench_kwargs=bench_kwargs)
+        run_results = _try_run(
+            model,
+            bench_fn,
+            bench_kwargs=bench_kwargs,
+            initial_batch_size=batch_size,
+            no_batch_size_retry=args.no_retry,
+        )
         if prefix and 'error' not in run_results:
             run_results = {'_'.join([prefix, k]): v for k, v in run_results.items()}
         model_results.update(run_results)

timm/data/__init__.py

@@ -6,7 +6,8 @@
 from .dataset_factory import create_dataset
 from .loader import create_loader
 from .mixup import Mixup, FastCollateMixup
-from .parsers import create_parser
+from .parsers import create_parser,\
+    get_img_extensions, is_img_extension, set_img_extensions, add_img_extensions, del_img_extensions
 from .real_labels import RealLabelsImagenet
 from .transforms import *
-from .transforms_factory import create_transform
+from .transforms_factory import create_transform

timm/data/config.py

@@ -64,11 +64,15 @@ def resolve_data_config(args, default_cfg={}, model=None, use_test_size=False, v
         new_config['std'] = default_cfg['std']
 
     # resolve default crop percentage
-    new_config['crop_pct'] = DEFAULT_CROP_PCT
+    crop_pct = DEFAULT_CROP_PCT
     if 'crop_pct' in args and args['crop_pct'] is not None:
-        new_config['crop_pct'] = args['crop_pct']
-    elif 'crop_pct' in default_cfg:
-        new_config['crop_pct'] = default_cfg['crop_pct']
+        crop_pct = args['crop_pct']
+    else:
+        if use_test_size and 'test_crop_pct' in default_cfg:
+            crop_pct = default_cfg['test_crop_pct']
+        elif 'crop_pct' in default_cfg:
+            crop_pct = default_cfg['crop_pct']
+    new_config['crop_pct'] = crop_pct
 
     if verbose:
         _logger.info('Data processing configuration for current model + dataset:')

timm/data/dataset_factory.py

@@ -26,8 +26,8 @@
     kmnist=KMNIST,
     fashion_mnist=FashionMNIST,
 )
-_TRAIN_SYNONYM = {'train', 'training'}
-_EVAL_SYNONYM = {'val', 'valid', 'validation', 'eval', 'evaluation'}
+_TRAIN_SYNONYM = dict(train=None, training=None)
+_EVAL_SYNONYM = dict(val=None, valid=None, validation=None, eval=None, evaluation=None)
 
 
 def _search_split(root, split):

timm/data/parsers/__init__.py

@@ -1 +1,2 @@
 from .parser_factory import create_parser
+from .img_extensions import *

timm/data/parsers/constants.py

@@ -0,0 +1,50 @@
+from copy import deepcopy
+
+__all__ = ['get_img_extensions', 'is_img_extension', 'set_img_extensions', 'add_img_extensions', 'del_img_extensions']
+
+
+IMG_EXTENSIONS = ('.png', '.jpg', '.jpeg')  # singleton, kept public for bwd compat use
+_IMG_EXTENSIONS_SET = set(IMG_EXTENSIONS)  # set version, private, kept in sync
+
+
+def _set_extensions(extensions):
+    global IMG_EXTENSIONS
+    global _IMG_EXTENSIONS_SET
+    dedupe = set()  # NOTE de-duping tuple while keeping original order
+    IMG_EXTENSIONS = tuple(x for x in extensions if x not in dedupe and not dedupe.add(x))
+    _IMG_EXTENSIONS_SET = set(extensions)
+
+
+def _valid_extension(x: str):
+    return x and isinstance(x, str) and len(x) >= 2 and x.startswith('.')
+
+
+def is_img_extension(ext):
+    return ext in _IMG_EXTENSIONS_SET
+
+
+def get_img_extensions(as_set=False):
+    return deepcopy(_IMG_EXTENSIONS_SET if as_set else IMG_EXTENSIONS)
+
+
+def set_img_extensions(extensions):
+    assert len(extensions)
+    for x in extensions:
+        assert _valid_extension(x)
+    _set_extensions(extensions)
+
+
+def add_img_extensions(ext):
+    if not isinstance(ext, (list, tuple, set)):
+        ext = (ext,)
+    for x in ext:
+        assert _valid_extension(x)
+    extensions = IMG_EXTENSIONS + tuple(ext)
+    _set_extensions(extensions)
+
+
+def del_img_extensions(ext):
+    if not isinstance(ext, (list, tuple, set)):
+        ext = (ext,)
+    extensions = tuple(x for x in IMG_EXTENSIONS if x not in ext)
+    _set_extensions(extensions)

timm/data/parsers/img_extensions.py

@@ -1,7 +1,6 @@
 import os
 
 from .parser_image_folder import ParserImageFolder
-from .parser_image_tar import ParserImageTar
 from .parser_image_in_tar import ParserImageInTar
 
 

timm/data/parsers/parser_factory.py

@@ -6,15 +6,35 @@
 Hacked together by / Copyright 2020 Ross Wightman
 """
 import os
+from typing import Dict, List, Optional, Set, Tuple, Union
 
 from timm.utils.misc import natural_key
 
-from .parser import Parser
 from .class_map import load_class_map
-from .constants import IMG_EXTENSIONS
+from .img_extensions import get_img_extensions
+from .parser import Parser
+
+
+def find_images_and_targets(
+        folder: str,
+        types: Optional[Union[List, Tuple, Set]] = None,
+        class_to_idx: Optional[Dict] = None,
+        leaf_name_only: bool = True,
+        sort: bool = True
+):
+    """ Walk folder recursively to discover images and map them to classes by folder names.
 
+    Args:
+        folder: root of folder to recrusively search
+        types: types (file extensions) to search for in path
+        class_to_idx: specify mapping for class (folder name) to class index if set
+        leaf_name_only: use only leaf-name of folder walk for class names
+        sort: re-sort found images by name (for consistent ordering)
 
-def find_images_and_targets(folder, types=IMG_EXTENSIONS, class_to_idx=None, leaf_name_only=True, sort=True):
+    Returns:
+        A list of image and target tuples, class_to_idx mapping
+    """
+    types = get_img_extensions(as_set=True) if not types else set(types)
     labels = []
     filenames = []
     for root, subdirs, files in os.walk(folder, topdown=False, followlinks=True):
@@ -51,7 +71,8 @@ def __init__(
         self.samples, self.class_to_idx = find_images_and_targets(root, class_to_idx=class_to_idx)
         if len(self.samples) == 0:
             raise RuntimeError(
-                f'Found 0 images in subfolders of {root}. Supported image extensions are {", ".join(IMG_EXTENSIONS)}')
+                f'Found 0 images in subfolders of {root}. '
+                f'Supported image extensions are {", ".join(get_img_extensions())}')
 
     def __getitem__(self, index):
         path, target = self.samples[index]

timm/data/parsers/parser_image_folder.py

@@ -9,20 +9,20 @@
 
 Hacked together by / Copyright 2020 Ross Wightman
 """
+import logging
 import os
-import tarfile
 import pickle
-import logging
-import numpy as np
+import tarfile
 from glob import glob
-from typing import List, Dict
+from typing import List, Tuple, Dict, Set, Optional, Union
+
+import numpy as np
 
 from timm.utils.misc import natural_key
 
-from .parser import Parser
 from .class_map import load_class_map
-from .constants import IMG_EXTENSIONS
-
+from .img_extensions import get_img_extensions
+from .parser import Parser
 
 _logger = logging.getLogger(__name__)
 CACHE_FILENAME_SUFFIX = '_tarinfos.pickle'
@@ -39,7 +39,7 @@ def reset(self):
         self.tf = None
 
 
-def _extract_tarinfo(tf: tarfile.TarFile, parent_info: Dict, extensions=IMG_EXTENSIONS):
+def _extract_tarinfo(tf: tarfile.TarFile, parent_info: Dict, extensions: Set[str]):
     sample_count = 0
     for i, ti in enumerate(tf):
         if not ti.isfile():
@@ -60,7 +60,14 @@ def _extract_tarinfo(tf: tarfile.TarFile, parent_info: Dict, extensions=IMG_EXTE
     return sample_count
 
 
-def extract_tarinfos(root, class_name_to_idx=None, cache_tarinfo=None, extensions=IMG_EXTENSIONS, sort=True):
+def extract_tarinfos(
+        root,
+        class_name_to_idx: Optional[Dict] = None,
+        cache_tarinfo: Optional[bool] = None,
+        extensions: Optional[Union[List, Tuple, Set]] = None,
+        sort: bool = True
+):
+    extensions = get_img_extensions(as_set=True) if not extensions else set(extensions)
     root_is_tar = False
     if os.path.isfile(root):
         assert os.path.splitext(root)[-1].lower() == '.tar'
@@ -176,8 +183,8 @@ def __init__(self, root, class_map='', cache_tarfiles=True, cache_tarinfo=None):
         self.samples, self.targets, self.class_name_to_idx, tarfiles = extract_tarinfos(
             self.root,
             class_name_to_idx=class_name_to_idx,
-            cache_tarinfo=cache_tarinfo,
-            extensions=IMG_EXTENSIONS)
+            cache_tarinfo=cache_tarinfo
+        )
         self.class_idx_to_name = {v: k for k, v in self.class_name_to_idx.items()}
         if len(tarfiles) == 1 and tarfiles[0][0] is None:
             self.root_is_tar = True