Add dd factory kwargs to eva, resnet

Author: rwightman

timm/models/eva.py

@@ -121,6 +121,8 @@ def __init__(
             qk_norm: bool = False,
             scale_norm: bool = True,
             rotate_half: bool = False,
+            device=None,
+            dtype=None,
     ):
         """
         Args:
@@ -139,6 +141,7 @@ def __init__(
             scale_norm: Enable normalization (scaling) of attention output with norm_layer
             rotate_half: Use half rotation layout instead of interleaved
         """
+        dd = {'device': device, 'dtype': dtype}
         super().__init__()
         if scale_norm or qk_norm:
             assert norm_layer is not None, 'norm_layer must be provided if qk_norm or scale_norm is True'
@@ -154,25 +157,25 @@ def __init__(
         self.rotate_half = rotate_half
 
         if qkv_fused:
-            self.qkv = nn.Linear(dim, attn_dim * 3, bias=False)
+            self.qkv = nn.Linear(dim, attn_dim * 3, bias=False, **dd)
             self.q_proj = self.k_proj = self.v_proj = None
             if qkv_bias:
-                self.q_bias = nn.Parameter(torch.zeros(attn_dim))
-                self.register_buffer('k_bias', torch.zeros(attn_dim), persistent=False)
-                self.v_bias = nn.Parameter(torch.zeros(attn_dim))
+                self.q_bias = nn.Parameter(torch.zeros(attn_dim, **dd))
+                self.register_buffer('k_bias', torch.zeros(attn_dim, **dd), persistent=False)
+                self.v_bias = nn.Parameter(torch.zeros(attn_dim, **dd))
             else:
                 self.q_bias = self.k_bias = self.v_bias = None
         else:
-            self.q_proj = nn.Linear(dim, attn_dim, bias=qkv_bias)
-            self.k_proj = nn.Linear(dim, attn_dim, bias=False)
-            self.v_proj = nn.Linear(dim, attn_dim, bias=qkv_bias)
+            self.q_proj = nn.Linear(dim, attn_dim, bias=qkv_bias, **dd)
+            self.k_proj = nn.Linear(dim, attn_dim, bias=False, **dd)
+            self.v_proj = nn.Linear(dim, attn_dim, bias=qkv_bias, **dd)
             self.qkv = None
             self.q_bias = self.k_bias = self.v_bias = None
-        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
-        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.q_norm = norm_layer(self.head_dim, **dd) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim, **dd) if qk_norm else nn.Identity()
         self.attn_drop = nn.Dropout(attn_drop)
-        self.norm = norm_layer(attn_dim) if scale_norm else nn.Identity()
-        self.proj = nn.Linear(attn_dim, dim)
+        self.norm = norm_layer(attn_dim, **dd) if scale_norm else nn.Identity()
+        self.proj = nn.Linear(attn_dim, dim, **dd)
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(
@@ -263,6 +266,8 @@ def __init__(
             act_layer: Callable = nn.GELU,
             norm_layer: Callable = LayerNorm,
             attn_head_dim: Optional[int] = None,
+            device=None,
+            dtype=None,
             **kwargs,
     ):
         """ Initialize the EVA transformer block.
@@ -286,8 +291,10 @@ def __init__(
             norm_layer: Normalization layer constructor
             attn_head_dim: Dimension of each attention head (if None, computed as dim // num_heads)
         """
+        dd = {'device': device, 'dtype': dtype}
         super().__init__()
-        self.norm1 = norm_layer(dim)
+
+        self.norm1 = norm_layer(dim, **dd)
         attn_cls = AttentionRope if attn_type == 'rope' else EvaAttention
         self.attn = attn_cls(
             dim,
@@ -301,11 +308,12 @@ def __init__(
             norm_layer=norm_layer,
             scale_norm=scale_attn_inner,
             rotate_half=rotate_half,
+            **dd,
         )
-        self.gamma_1 = nn.Parameter(init_values * torch.ones(dim)) if init_values is not None else None
+        self.gamma_1 = nn.Parameter(init_values * torch.ones(dim, **dd)) if init_values is not None else None
         self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
-        self.norm2 = norm_layer(dim)
+        self.norm2 = norm_layer(dim, **dd)
         hidden_features = int(dim * mlp_ratio)
         if swiglu_mlp:
             if scale_mlp or swiglu_align_to:
@@ -316,6 +324,7 @@ def __init__(
                     norm_layer=norm_layer if scale_mlp else None,
                     drop=proj_drop,
                     align_to=swiglu_align_to,
+                    **dd,
                 )
             else:
                 # w/o any extra norm, an impl with packed weights is used
@@ -326,6 +335,7 @@ def __init__(
                     act_layer=nn.SiLU,
                     gate_last=False,
                     drop=proj_drop,
+                    **dd,
                 )
         else:
             self.mlp = Mlp(
@@ -334,8 +344,9 @@ def __init__(
                 act_layer=act_layer,
                 norm_layer=norm_layer if scale_mlp else None,
                 drop=proj_drop,
+                **dd,
             )
-        self.gamma_2 = nn.Parameter(init_values * torch.ones(dim)) if init_values is not None else None
+        self.gamma_2 = nn.Parameter(init_values * torch.ones(dim, **dd)) if init_values is not None else None
         self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
     def forward(
@@ -376,6 +387,8 @@ def __init__(
             act_layer: Callable = nn.GELU,
             norm_layer: Callable = nn.LayerNorm,
             attn_head_dim: Optional[int] = None,
+            device=None,
+            dtype=None,
     ):
         """ Initialize the post-norm EVA transformer block.
 
@@ -398,7 +411,9 @@ def __init__(
             norm_layer: Normalization layer constructor
             attn_head_dim: Dimension of each attention head (if None, computed as dim // num_heads)
         """
+        dd = {'device': device, 'dtype': dtype}
         super().__init__()
+
         attn_cls = AttentionRope if attn_type == 'rope' else EvaAttention
         self.attn = attn_cls(
             dim,
@@ -412,8 +427,9 @@ def __init__(
             norm_layer=norm_layer,
             scale_norm=scale_attn_inner,
             rotate_half=rotate_half,
+            **dd,
         )
-        self.norm1 = norm_layer(dim)
+        self.norm1 = norm_layer(dim, **dd)
         self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
         hidden_features = int(dim * mlp_ratio)
@@ -426,6 +442,7 @@ def __init__(
                     norm_layer=norm_layer if scale_mlp else None,
                     drop=proj_drop,
                     align_to=swiglu_align_to,
+                    **dd,
                 )
             else:
                 # w/o any extra norm, an impl with packed fc1 weights is used, matches existing GluMLP
@@ -436,6 +453,7 @@ def __init__(
                     act_layer=nn.SiLU,
                     gate_last=False,
                     drop=proj_drop,
+                    **dd,
                 )
         else:
             self.mlp = Mlp(
@@ -444,8 +462,9 @@ def __init__(
                 act_layer=act_layer,
                 norm_layer=norm_layer if scale_mlp else None,
                 drop=proj_drop,
+                **dd,
             )
-        self.norm2 = norm_layer(dim)
+        self.norm2 = norm_layer(dim, **dd)
         self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
     def forward(
@@ -513,6 +532,8 @@ def __init__(
             dynamic_img_pad: bool = False,
             ref_feat_shape: Optional[Union[Tuple[int, int], int]] = None,
             head_init_scale: float = 0.001,
+            device=None,
+            dtype=None,
     ):
         """Initialize the EVA Vision Transformer model.
 
@@ -562,6 +583,7 @@ def __init__(
             head_init_scale: Initialization scale for classification head weights
         """
         super().__init__()
+        dd = {'device': device, 'dtype': dtype}
         assert global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map')
         self.num_classes = num_classes
         self.global_pool = global_pool
@@ -594,16 +616,17 @@ def __init__(
             dynamic_img_pad=dynamic_img_pad,
             bias=not use_pre_transformer_norm,
             **embed_args,
+            **dd,
         )
         num_patches = self.patch_embed.num_patches
         r = self.patch_embed.feat_ratio() if hasattr(self.patch_embed, 'feat_ratio') else patch_size
 
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None
-        self.reg_token = nn.Parameter(torch.zeros(1, num_reg_tokens, embed_dim)) if num_reg_tokens else None
+        self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim, **dd)) if class_token else None
+        self.reg_token = nn.Parameter(torch.empty(1, num_reg_tokens, embed_dim, **dd)) if num_reg_tokens else None
         self.cls_embed = class_token and self.reg_token is None
 
         num_pos_tokens = num_patches if no_embed_class else num_patches + self.num_prefix_tokens
-        self.pos_embed = nn.Parameter(torch.zeros(1, num_pos_tokens, embed_dim)) if use_abs_pos_emb else None
+        self.pos_embed = nn.Parameter(torch.empty(1, num_pos_tokens, embed_dim, **dd)) if use_abs_pos_emb else None
         self.pos_drop = nn.Dropout(p=pos_drop_rate)
         if patch_drop_rate > 0:
             self.patch_drop = PatchDropoutWithIndices(patch_drop_rate, num_prefix_tokens=self.num_prefix_tokens)
@@ -621,6 +644,7 @@ def __init__(
                 feat_shape=None if dynamic_img_size else self.patch_embed.grid_size,
                 temperature=rope_temperature,
                 grid_indexing=rope_grid_indexing,
+                **dd,
             )
             if rope_type == 'mixed':
                 rope_kwargs.update(dict(depth=depth))
@@ -636,7 +660,7 @@ def __init__(
         else:
             self.rope = None
 
-        self.norm_pre = norm_layer(embed_dim) if activate_pre_norm else nn.Identity()
+        self.norm_pre = norm_layer(embed_dim, **dd) if activate_pre_norm else nn.Identity()
 
         dpr = calculate_drop_path_rates(drop_path_rate, depth)  # stochastic depth decay rule
         block_fn = EvaBlockPostNorm if use_post_norm else EvaBlock
@@ -659,12 +683,13 @@ def __init__(
                 drop_path=dpr[i],
                 norm_layer=norm_layer,
                 init_values=init_values,
+                **dd,
             )
             for i in range(depth)])
         self.feature_info = [
             dict(module=f'blocks.{i}', num_chs=embed_dim, reduction=r) for i in range(depth)]
 
-        self.norm = norm_layer(embed_dim) if activate_post_norm else nn.Identity()
+        self.norm = norm_layer(embed_dim, **dd) if activate_post_norm else nn.Identity()
 
         if global_pool == 'map':
             self.attn_pool = AttentionPoolLatent(
@@ -673,23 +698,26 @@ def __init__(
                 mlp_ratio=attn_pool_mlp_ratio or mlp_ratio,
                 norm_layer=norm_layer,
                 act_layer=nn.GELU,
+                **dd,
             )
         else:
             self.attn_pool = None
-        self.fc_norm = norm_layer(embed_dim) if activate_fc_norm else nn.Identity()
+        self.fc_norm = norm_layer(embed_dim, **dd) if activate_fc_norm else nn.Identity()
         self.head_drop = nn.Dropout(drop_rate)
-        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        self.head = nn.Linear(embed_dim, num_classes, **dd) if num_classes > 0 else nn.Identity()
 
+        self.init_weights(head_init_scale=head_init_scale)
+
+    def init_weights(self, head_init_scale=None):
         self.apply(self._init_weights)
         if self.pos_embed is not None:
             trunc_normal_(self.pos_embed, std=.02)
         if self.cls_token is not None:
             trunc_normal_(self.cls_token, std=.02)
         if self.reg_token is not None:
             trunc_normal_(self.reg_token, std=.02)
-
         self.fix_init_weight()
-        if isinstance(self.head, nn.Linear):
+        if head_init_scale and isinstance(self.head, nn.Linear):
             trunc_normal_(self.head.weight, std=.02)
             self.head.weight.data.mul_(head_init_scale)
             self.head.bias.data.mul_(head_init_scale)