Pass through --model-kwargs (and --opt-kwargs for train) from command line through to model __init__. Update some models to improve arg overlay. Cleanup along the way.

Author: rwightman

benchmark.py

@@ -22,7 +22,7 @@
 from timm.layers import set_fast_norm
 from timm.models import create_model, is_model, list_models
 from timm.optim import create_optimizer_v2
-from timm.utils import setup_default_logging, set_jit_fuser, decay_batch_step, check_batch_size_retry
+from timm.utils import setup_default_logging, set_jit_fuser, decay_batch_step, check_batch_size_retry, ParseKwargs
 
 has_apex = False
 try:
@@ -108,12 +108,15 @@
                     help='Enable gradient checkpointing through model blocks/stages')
 parser.add_argument('--amp', action='store_true', default=False,
                     help='use PyTorch Native AMP for mixed precision training. Overrides --precision arg.')
+parser.add_argument('--amp-dtype', default='float16', type=str,
+                    help='lower precision AMP dtype (default: float16). Overrides --precision arg if args.amp True.')
 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')
+parser.add_argument('--model-kwargs', nargs='*', default={}, action=ParseKwargs)
 
 # codegen (model compilation) options
 scripting_group = parser.add_mutually_exclusive_group()
@@ -124,7 +127,6 @@
 scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
                              help="Enable AOT Autograd optimization.")
 
-
 # train optimizer parameters
 parser.add_argument('--opt', default='sgd', type=str, metavar='OPTIMIZER',
                     help='Optimizer (default: "sgd"')
@@ -168,19 +170,21 @@ def count_params(model: nn.Module):
 
 
 def resolve_precision(precision: str):
-    assert precision in ('amp', 'float16', 'bfloat16', 'float32')
-    use_amp = False
+    assert precision in ('amp', 'amp_bfloat16', 'float16', 'bfloat16', 'float32')
+    amp_dtype = None  # amp disabled
     model_dtype = torch.float32
     data_dtype = torch.float32
     if precision == 'amp':
-        use_amp = True
+        amp_dtype = torch.float16
+    elif precision == 'amp_bfloat16':
+        amp_dtype = torch.bfloat16
     elif precision == 'float16':
         model_dtype = torch.float16
         data_dtype = torch.float16
     elif precision == 'bfloat16':
         model_dtype = torch.bfloat16
         data_dtype = torch.bfloat16
-    return use_amp, model_dtype, data_dtype
+    return amp_dtype, model_dtype, data_dtype
 
 
 def profile_deepspeed(model, input_size=(3, 224, 224), batch_size=1, detailed=False):
@@ -228,9 +232,12 @@ def __init__(
         self.model_name = model_name
         self.detail = detail
         self.device = device
-        self.use_amp, self.model_dtype, self.data_dtype = resolve_precision(precision)
+        self.amp_dtype, self.model_dtype, self.data_dtype = resolve_precision(precision)
         self.channels_last = kwargs.pop('channels_last', False)
-        self.amp_autocast = partial(torch.cuda.amp.autocast, dtype=torch.float16) if self.use_amp else suppress
+        if self.amp_dtype is not None:
+            self.amp_autocast = partial(torch.cuda.amp.autocast, dtype=self.amp_dtype)
+        else:
+            self.amp_autocast = suppress
 
         if fuser:
             set_jit_fuser(fuser)
@@ -243,6 +250,7 @@ def __init__(
             drop_rate=kwargs.pop('drop', 0.),
             drop_path_rate=kwargs.pop('drop_path', None),
             drop_block_rate=kwargs.pop('drop_block', None),
+            **kwargs.pop('model_kwargs', {}),
         )
         self.model.to(
             device=self.device,
@@ -560,7 +568,7 @@ def _try_run(
 def benchmark(args):
     if args.amp:
         _logger.warning("Overriding precision to 'amp' since --amp flag set.")
-        args.precision = 'amp'
+        args.precision = 'amp' if args.amp_dtype == 'float16' else '_'.join(['amp', args.amp_dtype])
     _logger.info(f'Benchmarking in {args.precision} precision. '
                  f'{"NHWC" if args.channels_last else "NCHW"} layout. '
                  f'torchscript {"enabled" if args.torchscript else "disabled"}')

inference.py

@@ -20,7 +20,7 @@
 from timm.data import create_dataset, create_loader, resolve_data_config
 from timm.layers import apply_test_time_pool
 from timm.models import create_model
-from timm.utils import AverageMeter, setup_default_logging, set_jit_fuser
+from timm.utils import AverageMeter, setup_default_logging, set_jit_fuser, ParseKwargs
 
 try:
     from apex import amp
@@ -72,6 +72,8 @@
                     metavar='N', help='mini-batch size (default: 256)')
 parser.add_argument('--img-size', default=None, type=int,
                     metavar='N', help='Input image dimension, uses model default if empty')
+parser.add_argument('--in-chans', type=int, default=None, metavar='N',
+                    help='Image input channels (default: None => 3)')
 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,
@@ -110,6 +112,7 @@
                     help='lower precision AMP dtype (default: float16)')
 parser.add_argument('--fuser', default='', type=str,
                     help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
+parser.add_argument('--model-kwargs', nargs='*', default={}, action=ParseKwargs)
 
 scripting_group = parser.add_mutually_exclusive_group()
 scripting_group.add_argument('--torchscript', default=False, action='store_true',
@@ -170,12 +173,19 @@ def main():
         set_jit_fuser(args.fuser)
 
     # create model
+    in_chans = 3
+    if args.in_chans is not None:
+        in_chans = args.in_chans
+    elif args.input_size is not None:
+        in_chans = args.input_size[0]
+
     model = create_model(
         args.model,
         num_classes=args.num_classes,
-        in_chans=3,
+        in_chans=in_chans,
         pretrained=args.pretrained,
         checkpoint_path=args.checkpoint,
+        **args.model_kwargs,
     )
     if args.num_classes is None:
         assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'

timm/models/byobnet.py

@@ -218,7 +218,10 @@ def _rep_vgg_bcfg(d=(4, 6, 16, 1), wf=(1., 1., 1., 1.), groups=0):
 
 
 def interleave_blocks(
-        types: Tuple[str, str], d, every: Union[int, List[int]] = 1, first: bool = False, **kwargs
+        types: Tuple[str, str], d,
+        every: Union[int, List[int]] = 1,
+        first: bool = False,
+        **kwargs,
 ) -> Tuple[ByoBlockCfg]:
     """ interleave 2 block types in stack
     """
@@ -1587,15 +1590,32 @@ def __init__(
             in_chans=3,
             global_pool='avg',
             output_stride=32,
-            zero_init_last=True,
             img_size=None,
             drop_rate=0.,
             drop_path_rate=0.,
+            zero_init_last=True,
+            **kwargs,
     ):
+        """
+
+        Args:
+            cfg (ByoModelCfg): Model architecture configuration
+            num_classes (int): Number of classifier classes (default: 1000)
+            in_chans (int): Number of input channels (default: 3)
+            global_pool (str): Global pooling type (default: 'avg')
+            output_stride (int): Output stride of network, one of (8, 16, 32) (default: 32)
+            img_size (Union[int, Tuple[int]): Image size for fixed image size models (i.e. self-attn)
+            drop_rate (float): Dropout rate (default: 0.)
+            drop_path_rate (float): Stochastic depth drop-path rate (default: 0.)
+            zero_init_last (bool): Zero-init last weight of residual path
+            kwargs (dict): Extra kwargs overlayed onto cfg
+        """
         super().__init__()
         self.num_classes = num_classes
         self.drop_rate = drop_rate
         self.grad_checkpointing = False
+
+        cfg = replace(cfg, **kwargs)  # overlay kwargs onto cfg
         layers = get_layer_fns(cfg)
         if cfg.fixed_input_size:
             assert img_size is not None, 'img_size argument is required for fixed input size model'

timm/models/convnext.py

@@ -167,7 +167,7 @@ def __init__(
                 conv_bias=conv_bias,
                 use_grn=use_grn,
                 act_layer=act_layer,
-                norm_layer=norm_layer if conv_mlp else norm_layer_cl
+                norm_layer=norm_layer if conv_mlp else norm_layer_cl,
             ))
             in_chs = out_chs
         self.blocks = nn.Sequential(*stage_blocks)
@@ -184,16 +184,6 @@ def forward(self, x):
 class ConvNeXt(nn.Module):
     r""" ConvNeXt
         A PyTorch impl of : `A ConvNet for the 2020s`  - https://arxiv.org/pdf/2201.03545.pdf
-
-    Args:
-        in_chans (int): Number of input image channels. Default: 3
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (tuple(int)): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        drop_rate (float): Head dropout rate
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        ls_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
     """
 
     def __init__(
@@ -218,6 +208,28 @@ def __init__(
             drop_rate=0.,
             drop_path_rate=0.,
     ):
+        """
+        Args:
+            in_chans (int): Number of input image channels (default: 3)
+            num_classes (int): Number of classes for classification head (default: 1000)
+            global_pool (str): Global pooling type (default: 'avg')
+            output_stride (int): Output stride of network, one of (8, 16, 32) (default: 32)
+            depths (tuple(int)): Number of blocks at each stage. (default: [3, 3, 9, 3])
+            dims (tuple(int)): Feature dimension at each stage. (default: [96, 192, 384, 768])
+            kernel_sizes (Union[int, List[int]]: Depthwise convolution kernel-sizes for each stage (default: 7)
+            ls_init_value (float): Init value for Layer Scale (default: 1e-6)
+            stem_type (str): Type of stem (default: 'patch')
+            patch_size (int): Stem patch size for patch stem (default: 4)
+            head_init_scale (float): Init scaling value for classifier weights and biases (default: 1)
+            head_norm_first (bool): Apply normalization before global pool + head (default: False)
+            conv_mlp (bool): Use 1x1 conv in MLP, improves speed for small networks w/ chan last (default: False)
+            conv_bias (bool): Use bias layers w/ all convolutions (default: True)
+            use_grn (bool): Use Global Response Norm (ConvNeXt-V2) in MLP (default: False)
+            act_layer (Union[str, nn.Module]): Activation Layer
+            norm_layer (Union[str, nn.Module]): Normalization Layer
+            drop_rate (float): Head dropout rate (default: 0.)
+            drop_path_rate (float): Stochastic depth rate (default: 0.)
+        """
         super().__init__()
         assert output_stride in (8, 16, 32)
         kernel_sizes = to_ntuple(4)(kernel_sizes)
@@ -279,7 +291,7 @@ def __init__(
                 use_grn=use_grn,
                 act_layer=act_layer,
                 norm_layer=norm_layer,
-                norm_layer_cl=norm_layer_cl
+                norm_layer_cl=norm_layer_cl,
             ))
             prev_chs = out_chs
             # NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2

timm/models/cspnet.py

@@ -12,7 +12,7 @@
 
 Hacked together by / Copyright 2020 Ross Wightman
 """
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass, asdict, replace
 from functools import partial
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -518,7 +518,7 @@ def __init__(
             cross_linear=False,
             block_dpr=None,
             block_fn=BottleneckBlock,
-            **block_kwargs
+            **block_kwargs,
     ):
         super(CrossStage, self).__init__()
         first_dilation = first_dilation or dilation
@@ -558,7 +558,7 @@ def __init__(
                 bottle_ratio=bottle_ratio,
                 groups=groups,
                 drop_path=block_dpr[i] if block_dpr is not None else 0.,
-                **block_kwargs
+                **block_kwargs,
             ))
             prev_chs = block_out_chs
 
@@ -597,7 +597,7 @@ def __init__(
             cross_linear=False,
             block_dpr=None,
             block_fn=BottleneckBlock,
-            **block_kwargs
+            **block_kwargs,
     ):
         super(CrossStage3, self).__init__()
         first_dilation = first_dilation or dilation
@@ -635,7 +635,7 @@ def __init__(
                 bottle_ratio=bottle_ratio,
                 groups=groups,
                 drop_path=block_dpr[i] if block_dpr is not None else 0.,
-                **block_kwargs
+                **block_kwargs,
             ))
             prev_chs = block_out_chs
 
@@ -668,7 +668,7 @@ def __init__(
             avg_down=False,
             block_fn=BottleneckBlock,
             block_dpr=None,
-            **block_kwargs
+            **block_kwargs,
     ):
         super(DarkStage, self).__init__()
         first_dilation = first_dilation or dilation
@@ -715,7 +715,7 @@ def create_csp_stem(
         padding='',
         act_layer=nn.ReLU,
         norm_layer=nn.BatchNorm2d,
-        aa_layer=None
+        aa_layer=None,
 ):
     stem = nn.Sequential()
     feature_info = []
@@ -738,7 +738,7 @@ def create_csp_stem(
             stride=conv_stride,
             padding=padding if i == 0 else '',
             act_layer=act_layer,
-            norm_layer=norm_layer
+            norm_layer=norm_layer,
         ))
         stem_stride *= conv_stride
         prev_chs = chs
@@ -800,7 +800,7 @@ def create_csp_stages(
         cfg: CspModelCfg,
         drop_path_rate: float,
         output_stride: int,
-        stem_feat: Dict[str, Any]
+        stem_feat: Dict[str, Any],
 ):
     cfg_dict = asdict(cfg.stages)
     num_stages = len(cfg.stages.depth)
@@ -868,12 +868,27 @@ def __init__(
             global_pool='avg',
             drop_rate=0.,
             drop_path_rate=0.,
-            zero_init_last=True
+            zero_init_last=True,
+            **kwargs,
     ):
+        """
+        Args:
+            cfg (CspModelCfg): Model architecture configuration
+            in_chans (int): Number of input channels (default: 3)
+            num_classes (int): Number of classifier classes (default: 1000)
+            output_stride (int): Output stride of network, one of (8, 16, 32) (default: 32)
+            global_pool (str): Global pooling type (default: 'avg')
+            drop_rate (float): Dropout rate (default: 0.)
+            drop_path_rate (float): Stochastic depth drop-path rate (default: 0.)
+            zero_init_last (bool): Zero-init last weight of residual path
+            kwargs (dict): Extra kwargs overlayed onto cfg
+        """
         super().__init__()
         self.num_classes = num_classes
         self.drop_rate = drop_rate
         assert output_stride in (8, 16, 32)
+
+        cfg = replace(cfg, **kwargs)  # overlay kwargs onto cfg
         layer_args = dict(
             act_layer=cfg.act_layer,
             norm_layer=cfg.norm_layer,