dd factory kwargs for fastvit, convnext, mambaout

Author: rwightman

timm/models/convnext.py

@@ -44,10 +44,27 @@
 import torch.nn as nn
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
-from timm.layers import trunc_normal_, AvgPool2dSame, DropPath, calculate_drop_path_rates, Mlp, GlobalResponseNormMlp, \
- LayerNorm2d, LayerNorm, RmsNorm2d, RmsNorm, create_conv2d, get_act_layer, get_norm_layer, make_divisible, to_ntuple
-from timm.layers import SimpleNorm2d, SimpleNorm
-from timm.layers import NormMlpClassifierHead, ClassifierHead
+from timm.layers import (
+ trunc_normal_,
+ AvgPool2dSame,
+ DropPath,
+ calculate_drop_path_rates,
+ Mlp,
+ GlobalResponseNormMlp,
+ LayerNorm2d,
+ LayerNorm,
+ RmsNorm2d,
+ RmsNorm,
+ SimpleNorm2d,
+ SimpleNorm,
+ create_conv2d,
+ get_act_layer,
+ get_norm_layer,
+ make_divisible,
+ to_ntuple,
+ NormMlpClassifierHead,
+ ClassifierHead,
+)
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
 from ._manipulate import named_apply, checkpoint_seq
@@ -59,7 +76,15 @@
 class Downsample(nn.Module):
  """Downsample module for ConvNeXt."""
 
- def __init__(self, in_chs: int, out_chs: int, stride: int = 1, dilation: int = 1) -> None:
+ def __init__(
+ self,
+ in_chs: int,
+ out_chs: int,
+ stride: int = 1,
+ dilation: int = 1,
+ device=None,
+ dtype=None,
+ ) -> None:
  """Initialize Downsample module.
 
  Args:
@@ -68,6 +93,7 @@ def __init__(self, in_chs: int, out_chs: int, stride: int = 1, dilation: int = 1
  stride: Stride for downsampling.
  dilation: Dilation rate.
  """
+ dd = {'device': device, 'dtype': dtype}
  super().__init__()
  avg_stride = stride if dilation == 1 else 1
  if stride > 1 or dilation > 1:
@@ -77,7 +103,7 @@ def __init__(self, in_chs: int, out_chs: int, stride: int = 1, dilation: int = 1
  self.pool = nn.Identity()
 
  if in_chs != out_chs:
- self.conv = create_conv2d(in_chs, out_chs, 1, stride=1)
+ self.conv = create_conv2d(in_chs, out_chs, 1, stride=1, **dd)
  else:
  self.conv = nn.Identity()
 
@@ -115,6 +141,8 @@ def __init__(
  act_layer: Union[str, Callable] = 'gelu',
  norm_layer: Optional[Callable] = None,
  drop_path: float = 0.,
+ device=None,
+ dtype=None,
  ):
  """
 
@@ -133,6 +161,7 @@ def __init__(
  norm_layer: Normalization layer (defaults to LN if not specified).
  drop_path: Stochastic depth probability.
  """
+ dd = {'device': device, 'dtype': dtype}
  super().__init__()
  out_chs = out_chs or in_chs
  dilation = to_ntuple(2)(dilation)
@@ -149,12 +178,18 @@ def __init__(
  dilation=dilation[0],
  depthwise=True,
  bias=conv_bias,
+ **dd,
+ )
+ self.norm = norm_layer(out_chs, **dd)
+ self.mlp = mlp_layer(
+ out_chs,
+ int(mlp_ratio * out_chs),
+ act_layer=act_layer,
+ **dd,
  )
- self.norm = norm_layer(out_chs)
- self.mlp = mlp_layer(out_chs, int(mlp_ratio * out_chs), act_layer=act_layer)
- self.gamma = nn.Parameter(ls_init_value * torch.ones(out_chs)) if ls_init_value is not None else None
+ self.gamma = nn.Parameter(ls_init_value * torch.ones(out_chs, **dd)) if ls_init_value is not None else None
  if in_chs != out_chs or stride != 1 or dilation[0] != dilation[1]:
- self.shortcut = Downsample(in_chs, out_chs, stride=stride, dilation=dilation[0])
+ self.shortcut = Downsample(in_chs, out_chs, stride=stride, dilation=dilation[0], **dd)
  else:
  self.shortcut = nn.Identity()
  self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
@@ -196,7 +231,9 @@ def __init__(
  use_grn: bool = False,
  act_layer: Union[str, Callable] = 'gelu',
  norm_layer: Optional[Callable] = None,
- norm_layer_cl: Optional[Callable] = None
+ norm_layer_cl: Optional[Callable] = None,
+ device=None,
+ dtype=None,
  ) -> None:
  """Initialize ConvNeXt stage.
 
@@ -216,14 +253,15 @@ def __init__(
  norm_layer: Normalization layer.
  norm_layer_cl: Normalization layer for channels last.
  """
+ dd = {'device': device, 'dtype': dtype}
  super().__init__()
  self.grad_checkpointing = False
 
  if in_chs != out_chs or stride > 1 or dilation[0] != dilation[1]:
  ds_ks = 2 if stride > 1 or dilation[0] != dilation[1] else 1
  pad = 'same' if dilation[1] > 1 else 0 # same padding needed if dilation used
  self.downsample = nn.Sequential(
- norm_layer(in_chs),
+ norm_layer(in_chs, **dd),
  create_conv2d(
  in_chs,
  out_chs,
@@ -232,6 +270,7 @@ def __init__(
  dilation=dilation[0],
  padding=pad,
  bias=conv_bias,
+ **dd,
  ),
  )
  in_chs = out_chs
@@ -253,6 +292,7 @@ def __init__(
  use_grn=use_grn,
  act_layer=act_layer,
  norm_layer=norm_layer if conv_mlp else norm_layer_cl,
+ **dd,
  ))
  in_chs = out_chs
  self.blocks = nn.Sequential(*stage_blocks)
@@ -324,6 +364,8 @@ def __init__(
  norm_eps: Optional[float] = None,
  drop_rate: float = 0.,
  drop_path_rate: float = 0.,
+ device=None,
+ dtype=None,
  ):
  """
  Args:
@@ -349,6 +391,7 @@ def __init__(
  drop_path_rate: Stochastic depth drop rate.
  """
  super().__init__()
+ dd = {'device': device, 'dtype': dtype}
  assert output_stride in (8, 16, 32)
  kernel_sizes = to_ntuple(4)(kernel_sizes)
  norm_layer, norm_layer_cl = _get_norm_layers(norm_layer, conv_mlp, norm_eps)
@@ -362,17 +405,17 @@ def __init__(
  if stem_type == 'patch':
  # NOTE: this stem is a minimal form of ViT PatchEmbed, as used in SwinTransformer w/ patch_size = 4
  self.stem = nn.Sequential(
- nn.Conv2d(in_chans, dims[0], kernel_size=patch_size, stride=patch_size, bias=conv_bias),
- norm_layer(dims[0]),
+ nn.Conv2d(in_chans, dims[0], kernel_size=patch_size, stride=patch_size, bias=conv_bias, **dd),
+ norm_layer(dims[0], **dd),
  )
  stem_stride = patch_size
  else:
  mid_chs = make_divisible(dims[0] // 2) if 'tiered' in stem_type else dims[0]
  self.stem = nn.Sequential(*filter(None, [
- nn.Conv2d(in_chans, mid_chs, kernel_size=3, stride=2, padding=1, bias=conv_bias),
+ nn.Conv2d(in_chans, mid_chs, kernel_size=3, stride=2, padding=1, bias=conv_bias, **dd),
  act_layer() if 'act' in stem_type else None,
- nn.Conv2d(mid_chs, dims[0], kernel_size=3, stride=2, padding=1, bias=conv_bias),
- norm_layer(dims[0]),
+ nn.Conv2d(mid_chs, dims[0], kernel_size=3, stride=2, padding=1, bias=conv_bias, **dd),
+ norm_layer(dims[0], **dd),
  ]))
  stem_stride = 4
 
@@ -406,6 +449,7 @@ def __init__(
  act_layer=act_layer,
  norm_layer=norm_layer,
  norm_layer_cl=norm_layer_cl,
+ **dd,
  ))
  prev_chs = out_chs
  # NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2
@@ -417,12 +461,13 @@ def __init__(
  # otherwise pool -> norm -> fc, the default ConvNeXt ordering (pretrained FB weights)
  if head_norm_first:
  assert not head_hidden_size
- self.norm_pre = norm_layer(self.num_features)
+ self.norm_pre = norm_layer(self.num_features, **dd)
  self.head = ClassifierHead(
  self.num_features,
  num_classes,
  pool_type=global_pool,
  drop_rate=self.drop_rate,
+ **dd,
  )
  else:
  self.norm_pre = nn.Identity()
@@ -434,6 +479,7 @@ def __init__(
  drop_rate=self.drop_rate,
  norm_layer=norm_layer,
  act_layer='gelu',
+ **dd,
  )
  self.head_hidden_size = self.head.num_features
  named_apply(partial(_init_weights, head_init_scale=head_init_scale), self)