Upload model

README.md

@@ -1,9 +1,9 @@
 ---
 library_name: transformers
-pipeline_tag: image-feature-extraction
 license: other
 license_name: nvidia-open-model-license
 license_link: https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf
+pipeline_tag: image-feature-extraction
 ---
 
 # Model Overview

adaptor_generic.py

@@ -19,9 +19,23 @@ class GenericAdaptor(AdaptorBase):
     def __init__(self, main_config: Namespace, adaptor_config, state, mlp_config=None):
         super().__init__()
 
+        extra_args = dict()
+        ups = None
+        ups_rank = None
+        if adaptor_config is not None:
+            ups = adaptor_config.get('fd_upsample_factor', None)
+            ups_rank = adaptor_config.get('fd_upsample_rank', None)
+        elif mlp_config is not None:
+            ups = mlp_config["feature"].get('upsample_factor', None)
+            ups_rank = mlp_config["feature"].get('upsample_rank', None)
+        if ups is not None:
+            extra_args['upsample_factor'] = ups
+            extra_args['upsample_rank'] = ups_rank
+
         if state is not None:
-            self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.')
-            self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.')
+            spectral_heads = getattr(main_config, 'spectral_heads', False)
+            self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.', spectral_weights=spectral_heads)
+            self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.', spectral_weights=spectral_heads, **extra_args)
         else:
             assert mlp_config is not None, "Config must not be None if state is None"
 
@@ -38,16 +52,17 @@ class GenericAdaptor(AdaptorBase):
                 mlp_config["feature"]["hidden_dim"],
                 mlp_config["feature"]["output_dim"],
                 mlp_config["feature"]["num_inner"],
+                **extra_args
             )
 
     def forward(self, input: AdaptorInput) -> RadioOutput:
         # Convert input'd type to the type of the first parameter of the adaptor.
         first_param = next(self.parameters())
         summary = self.head_mlp(input.summary.to(dtype=first_param.dtype)).to(dtype=input.summary.dtype)
-        feat = self.feat_mlp(input.features.to(dtype=first_param.dtype)).to(dtype=input.features.dtype)
+        feat = self.feat_mlp(input.features.to(dtype=first_param.dtype), images=input.images, patch_size=input.patch_size).to(dtype=input.features.dtype)
 
         if input.feature_fmt == 'NCHW':
-            feat = (feat.reshape(feat.shape[0], input.images.shape[-2] // input.patch_size, input.images.shape[-1] // input.patch_size, feat.shape[2])
+            feat = (feat.reshape(feat.shape[0], input.images.shape[-2] // input.patch_size * self.feat_mlp.upsample_factor, input.images.shape[-1] // input.patch_size * self.feat_mlp.upsample_factor, feat.shape[2])
                         .permute(0, 3, 1, 2)
             )
 

adaptor_mlp.py

@@ -6,7 +6,7 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 import math
-from typing import Dict
+from typing import Dict, Optional
 
 import torch
 from torch import nn
@@ -14,6 +14,8 @@ from torch import nn
 from einops import rearrange
 from timm.models.vision_transformer import Block
 
+from .enable_spectral_reparam import disable_spectral_reparam, enable_spectral_reparam
+
 
 class MLP(nn.Module):
     def __init__(self, input_size: int, hidden_size: int, output_size: int,
@@ -51,6 +53,8 @@ class MLP2(nn.Module):
                  num_inner: int = 0,
                  pre_norm: bool = False, device: torch.device = None,
                  upsample_factor: int = 1,
+                 upsample_rank: int = None,
+                 from_config: bool = False,
                  **kwargs):
         super().__init__()
 
@@ -60,10 +64,12 @@ class MLP2(nn.Module):
         ) if pre_norm else nn.Identity()
 
         self.upsample_factor = upsample_factor
-        self._real_output_dim = output_size
+        sq_ups = upsample_factor ** 2
+
+        self._real_output_dim = output_size // sq_ups
 
-        hidden_size *= upsample_factor
-        output_size *= (upsample_factor ** 2)
+        # hidden_size *= upsample_factor
+        # output_size *= (upsample_factor ** 2)
 
         self.fc1 = nn.Linear(input_size, hidden_size, device=device)
 
@@ -82,7 +88,7 @@ class MLP2(nn.Module):
             nn.Linear(hidden_size, output_size, device=device),
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, images: Optional[torch.Tensor] = None, patch_size: Optional[int] = None) -> torch.Tensor:
         x = self.pre_norm(x)
         x = self.fc1(x)
         for block in self.blocks:
@@ -90,8 +96,12 @@ class MLP2(nn.Module):
         x = self.final(x)
 
         if self.upsample_factor > 1:
-            h = w = int(math.sqrt(x.shape[1]))
-            x = rearrange(x, 'b (h w) (u1 u2 c) -> b (u1 h u2 w) c',
+            if images is None:
+                raise ValueError(f'`images` cannot be `None` when the head\'s `upsample_factor > 1`!')
+            if patch_size is None:
+                raise ValueError(f'`patch_size` cannot be `None` when the head\'s `upsample_factor > 1`!')
+            h, w = tuple(d // patch_size for d in images.shape[-2:])
+            x = rearrange(x, 'b (h w) (u1 u2 c) -> b (h u1 w u2) c',
                           h=h, w=w, u1=self.upsample_factor, u2=self.upsample_factor,
                           c=self._real_output_dim)
 
@@ -113,20 +123,22 @@ def strip_prefix(state: Dict[str, torch.Tensor], prefix: str):
     return state
 
 
-def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = ''):
+def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False):
     state = strip_prefix(state, prefix)
 
+    weight_suffix = 'weight' if not spectral_weights else 'parametrizations.weight.original'
+
     if version == 'v1':
-        hidden_dim, input_dim = state['fc1.weight'].shape
-        output_dim = state['fc2.weight'].shape[0]
+        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
+        output_dim = state[f'fc2.{weight_suffix}'].shape[0]
 
         for num_inner in range(1000):
             k = f'inner.{num_inner}.0.weight'
             if k not in state:
                 break
     elif version == 'v2':
-        hidden_dim, input_dim = state['fc1.weight'].shape
-        output_dim = state['final.2.weight'].shape[0]
+        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
+        output_dim = state[f'final.2.{weight_suffix}'].shape[0]
 
         for num_inner in range(1000):
             k = f'blocks.{num_inner}.0.weight'
@@ -138,19 +150,25 @@ def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix
     return input_dim, hidden_dim, output_dim, num_inner
 
 
-def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int, output_dim: int, num_inner: int):
-    ret: nn.Module = MLP_FACTORY[version](input_dim, hidden_dim, output_dim, num_inner)
+def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int, output_dim: int, num_inner: int, **kwargs):
+    ret: nn.Module = MLP_FACTORY[version](input_dim, hidden_dim, output_dim, num_inner, from_config=True, **kwargs)
 
     return ret
 
 
-def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = ''):
+def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False, **kwargs):
     state = strip_prefix(state, prefix)
 
-    input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state)
+    input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state, spectral_weights=spectral_weights)
+
+    ret: nn.Module = create_mlp_from_config(version, input_dim, hidden_dim, output_dim, num_inner, **kwargs)
 
-    ret: nn.Module = create_mlp_from_config(version, input_dim, hidden_dim, output_dim, num_inner)
+    if spectral_weights:
+        enable_spectral_reparam(ret, init_norm_to_current=False, state_dict_guidance=state)
 
     ret.load_state_dict(state)
 
+    if spectral_weights:
+        disable_spectral_reparam(ret)
+
     return ret

common.py

@@ -87,6 +87,22 @@ RESOURCE_MAP = {
         max_resolution=2048,
         preferred_resolution=Resolution(512, 512),
     ),
+    # C-RADIO
+    "c-radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/C-RADIOv2-g/resolve/main/c-radio_v2-g_half.pth.tar",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=8,
+    ),
+    "c-radio_v3-l": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-L/resolve/main/c-radio-v3_l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
 }
 
 DEFAULT_VERSION = "radio_v2.5-h"

dual_hybrid_vit.py

@@ -0,0 +1,213 @@
+from logging import getLogger
+from typing import Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from timm.models import register_model
+from timm.models import vision_transformer as tvit
+from timm.models import convnext as tconv
+
+from einops import rearrange
+
+from . import extra_timm_models as et
+
+
+class Fuser(nn.Module):
+    def __init__(self, src_dim: int, tgt_dim: int, gated: bool = True):
+        super().__init__()
+        self.gated = gated
+
+        mid_dim = max(src_dim, tgt_dim) * 2
+
+        self.fwd = nn.Sequential(
+            nn.Conv2d(src_dim, mid_dim, kernel_size=3, stride=1, padding=1),
+            nn.GELU(),
+            nn.Conv2d(mid_dim, tgt_dim * (2 if gated else 1), kernel_size=3, stride=1, padding=1),
+        )
+
+    def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
+        if src.ndim == 3:
+            shape = tgt.shape[-2:]
+        else:
+            shape = src.shape[-2:]
+
+        nd = shape[0] * shape[1]
+
+        if src.ndim == 3:
+            src = src[:, -nd:].reshape(src.shape[0], src.shape[2], *shape)
+
+        if tgt.ndim == 3:
+            tgt_pre = tgt[:, :-nd]
+            tgt = tgt[:, -nd:].reshape(tgt.shape[0], tgt.shape[2], *shape)
+        else:
+            tgt_pre = None
+
+        pred = self.fwd(src)
+
+        if self.gated:
+            g, pred = torch.chunk(pred, 2, dim=1)
+
+            g = F.sigmoid(g)
+
+            pred = g * pred
+
+        tgt = tgt + pred
+
+        if tgt_pre is not None:
+            tgt = rearrange(tgt, 'b c h w -> b (h w) c')
+            tgt = torch.cat([tgt_pre, tgt], dim=1)
+
+        return tgt
+
+
+class AttnDownsample(nn.Module):
+    def __init__(self, dim: int, window_size: int, num_heads: int = 16):
+        super().__init__()
+        self.q = nn.Parameter(torch.randn(1, num_heads, 1, dim // num_heads) * 0.01)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+        self.window_size = window_size
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+    def forward(self, x: torch.Tensor, twod_shape: Tuple[int, int]) -> torch.Tensor:
+        ntok = twod_shape[0] * twod_shape[1]
+        x_pre = x[:, :-ntok]
+
+        B = x.shape[0]
+        ds_hw = tuple(s // self.window_size for s in twod_shape)
+
+        x_spat = rearrange(
+            x[:, -ntok:],
+            'b (h d1 w d2) c -> (b h w) (d1 d2) c',
+            h=ds_hw[0], w=ds_hw[1],
+            d1=self.window_size, d2=self.window_size,
+        )
+
+        B, N, C = x_spat.shape
+
+        k, v = self.kv(x_spat).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+
+        q = (self.q * self.scale).expand(B, -1, -1, -1)
+        attn = q @ k.transpose(-2, -1)
+        attn = F.softmax(attn, dim=-1)
+        x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, C)
+        x = self.proj(x)
+
+        x = rearrange(x, '(b h w) c -> b (h w) c', b=x_pre.shape[0], h=ds_hw[0], w=ds_hw[1])
+
+        x = torch.cat([x_pre, x], dim=1)
+        return x
+
+
+class HybridModel(nn.Module):
+    def __init__(self, vit: tvit.VisionTransformer, conv: tconv.ConvNeXt, pretrained: bool = False,
+                 concatenate: bool = False, **kwargs):
+        super().__init__()
+        self.conv = conv
+        self.vit = vit
+        self.concatenate = concatenate
+
+        conv.stages = nn.ModuleList(conv.stages)
+        vit.blocks = nn.ModuleList(vit.blocks)
+
+        self._half_vit_idx = len(vit.blocks) // 2 + 1
+
+        self._half_conv_idx = None
+        x = torch.empty(1, 3, 256, 256)
+        x = self.conv.stem(x)
+        for i in range(len(conv.stages)):
+            x = conv.stages[i](x)
+            if self._half_conv_idx is None and x.shape[-2:] == (16, 16):
+                self._half_conv_idx = i + 1
+                half_conv_dim = x.shape[1]
+            final_conv_dim = x.shape[1]
+
+        self.vit_to_conv_fusion = Fuser(vit.embed_dim, half_conv_dim)
+        self.conv_to_vit_fusion = Fuser(half_conv_dim, vit.embed_dim)
+        self.vit_ds = AttnDownsample(vit.embed_dim, window_size=2)
+
+        embed_dim = vit.embed_dim + (final_conv_dim if concatenate else 0)
+        if not concatenate:
+            self.final_fuse = Fuser(final_conv_dim, vit.embed_dim, gated=False)
+        self.final_block = tvit.Block(embed_dim, num_heads=16)
+
+        self.embed_dim = embed_dim
+
+    @property
+    def patch_size(self):
+        return 32
+
+    @property
+    def no_fsdp_wrap_types(self):
+        return {tvit.VisionTransformer, tconv.ConvNeXt}
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_features(x)
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        y_vit = self.vit.patch_generator(x)
+
+        for i in range(self._half_vit_idx):
+            y_vit = self.vit.blocks[i](y_vit)
+
+        y_conv = self.conv.stem(x)
+        for i in range(self._half_conv_idx):
+            y_conv = self.conv.stages[i](y_conv)
+
+        y_vit, y_conv = self.conv_to_vit_fusion(y_conv, y_vit), self.vit_to_conv_fusion(y_vit, y_conv)
+
+        y_vit = self.vit_ds(y_vit, y_conv.shape[-2:])
+
+        for i in range(self._half_vit_idx, len(self.vit.blocks)):
+            y_vit = self.vit.blocks[i](y_vit)
+
+        for i in range(self._half_conv_idx, len(self.conv.stages)):
+            y_conv = self.conv.stages[i](y_conv)
+
+        if self.concatenate:
+            y_conv = rearrange(y_conv, 'b c h w -> b (h w) c')
+            # Average pool across the board, and replicate for each cls/register token
+            conv_summary = y_conv.mean(dim=1, keepdim=True).expand(-1, self.vit.patch_generator.num_cls_patches, -1)
+            y_conv = torch.cat([conv_summary, y_conv], dim=1)
+            y = torch.cat([y_vit, y_conv], dim=2)
+        else:
+            y = self.final_fuse(y_conv, y_vit)
+        y = self.final_block(y)
+
+        summary = y[:, :self.vit.patch_generator.num_cls_tokens]
+        features = y[:, self.vit.patch_generator.num_cls_patches:]
+
+        return summary, features
+
+
+@register_model
+def hybrid_base(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_base(pretrained=pretrained, **cfg)
+    vit = tvit.vit_base_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)
+
+
+@register_model
+def hybrid_large(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_large(pretrained=pretrained, **cfg)
+    vit = tvit.vit_large_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)
+
+
+@register_model
+def hybrid_huge(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_huge(pretrained=pretrained, **cfg)
+    vit = et.vit_huge_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)

enable_cpe_support.py

@@ -19,6 +19,7 @@ from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermedi
 from .extra_models import DinoWrapper
 from .vit_patch_generator import ViTPatchGenerator
 from .forward_intermediates import forward_intermediates
+from .dual_hybrid_vit import HybridModel
 
 
 def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor:
@@ -161,7 +162,9 @@ def enable_cpe(model: nn.Module,
 ):
     if isinstance(model, VisionTransformer):
         _enable_cpe_for_timm_vit(model, *args, **kwargs)
-    elif True: # isinstance(model, DinoWrapper):
+    elif isinstance(model, DinoWrapper):
         _enable_cpe_for_dv2_reg_vit(model, *args, **kwargs)
+    elif isinstance(model, HybridModel):
+        _enable_cpe_for_timm_vit(model.vit, *args, **kwargs)
     else:
         raise ValueError(f'CPE not supported for this model type: {type(model)}')

enable_spectral_reparam.py

@@ -155,7 +155,7 @@ def enable_spectral_reparam(model: Union[nn.Module, List[nn.Module]],
             return True
 
         p_name = f'{name}.parametrizations'
-        is_prm = any(k for k in state_dict_guidance if k.startswith(p_name))
+        is_prm = any(k for k in state_dict_guidance if k.startswith(p_name) and k.endswith('_sn_version'))
         return is_prm
 
     def parametrize_linear(linear: nn.Linear):

extra_timm_models.py

@@ -6,10 +6,24 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import math
+import warnings
+
+import torch
 from torch import nn
+from torch.nn import functional as F
 
 from timm.models import register_model
-from timm.models.vision_transformer import VisionTransformer, _create_vision_transformer, Mlp
+from timm.models.vision_transformer import (
+    VisionTransformer,
+    _create_vision_transformer as _timm_create_vision_transformer,
+    Mlp,
+    Block,
+    LayerScale as TIMMLayerScale,
+)
+
+# Import these to also register them
+from . import dinov2_arch
 
 
 @register_model
@@ -40,6 +54,34 @@ def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     return model
 
 
+@register_model
+def vit_base_patch16_v2_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5,
+        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
+    )
+    model = _create_vision_transformer(
+        'vit_base_patch14_reg4_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_large_patch16_v2_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    name = 'vit_large_patch14_reg4_dinov2'
+    model_args = dict(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, init_values=1e-5,
+        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
+    )
+    model = _create_vision_transformer(name, pretrained=pretrained, **dict(model_args, **kwargs))
+
+    return model
+
 @register_model
 def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
     """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
@@ -66,9 +108,99 @@ def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransforme
     return model
 
 
+@register_model
+def vit_giant_patch16_224(pretrained=False, scaled_ln: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-giant model (ViT-g/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model_args = dict(patch_size=16, embed_dim=1536, depth=40, num_heads=24)
+    model = _create_vision_transformer('vit_giant_patch16_224', pretrained=False, **dict(model_args, **kwargs))
+    if scaled_ln:
+        _apply_scaled_ln(model)
+    return model
+
+
 @register_model
 def vit_bigG_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
     model_args = dict(patch_size=14, embed_dim=1664, depth=48, num_heads=16, init_values=1e-6)
     model = _create_vision_transformer('vit_bigG_patch14', pretrained=False, **dict(model_args, **kwargs))
     return model
 
+
+def _create_vision_transformer(*args, **kwargs):
+    model = _timm_create_vision_transformer(*args, **kwargs)
+    _patch_layer_scale(model)
+    return model
+
+
+def _patch_layer_scale(model: VisionTransformer):
+    def replace_ls(old_ls: TIMMLayerScale):
+        new_ls = dinov2_arch.LayerScale(old_ls.gamma.shape[0], inplace=old_ls.inplace)
+        new_ls.load_state_dict(old_ls.state_dict())
+        return new_ls
+
+    # Monkey patch: Replace TIMM's LayerScale with our modified DINOv2 one, that uses a param name
+    # other than gamma, so that HFHub doesn't mess with it!
+    for mod in model.modules():
+        if isinstance(mod, Block):
+            if isinstance(mod.ls1, TIMMLayerScale):
+                mod.ls1 = replace_ls(mod.ls1)
+            if isinstance(mod.ls2, TIMMLayerScale):
+                mod.ls2 = replace_ls(mod.ls2)
+    pass
+
+
+class ScaledLayerNorm(nn.LayerNorm):
+    '''
+    https://arxiv.org/pdf/2502.05795v1
+    '''
+    def __init__(self, ln_base: nn.LayerNorm, depth: int = 0):
+        super().__init__(ln_base.normalized_shape, eps=ln_base.eps, elementwise_affine=ln_base.elementwise_affine)
+        self.load_state_dict(ln_base.state_dict())
+        self.register_buffer('ln_scale', torch.tensor(1.0 / math.sqrt(depth)), persistent=False)
+
+    def forward(self, x):
+        y = super().forward(x)
+        y = y * self.ln_scale
+        return y
+
+
+class DyT(nn.Module):
+    def __init__(self, C: int, init_alpha: float):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.full((1,), init_alpha))
+        self.gamma = nn.Parameter(torch.ones(C))
+        self.beta = nn.Parameter(torch.zeros(C))
+
+    def forward(self, x: torch.Tensor):
+        x = F.tanh(self.alpha * x)
+        return self.gamma * x + self.beta
+
+@register_model
+def vit_large_dyt_patch16_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16)
+    model = _create_vision_transformer('vit_large_dyt_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
+
+    def _replace_ln_with_dyt(ln: nn.LayerNorm, depth: int):
+        return DyT(ln.normalized_shape[0], init_alpha=0.9)
+    _replace_ln(model, _replace_ln_with_dyt)
+
+    return model
+
+
+def _apply_scaled_ln(model: VisionTransformer):
+    warnings.warn('Post-LayerNorm scaling activated!')
+
+    _replace_ln(model, lambda ln, depth: ScaledLayerNorm(ln, depth=depth))
+
+def _replace_ln(model: VisionTransformer, fn):
+    def _inner_replace_ln(block: Block, depth: int, key: str):
+        prev = getattr(block, key)
+        if isinstance(prev, nn.LayerNorm):
+            setattr(block, key, fn(prev, depth=depth))
+
+    for i, block in enumerate(model.blocks):
+        _inner_replace_ln(block, i + 1, 'norm1')
+        _inner_replace_ln(block, i + 1, 'norm2')

forward_intermediates.py

@@ -6,7 +6,7 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-from typing import Callable, List, Optional, Set, Tuple, Union, Any, Iterable
+from typing import Callable, Dict, List, Optional, Set, Tuple, Union, Any, Iterable
 from types import MethodType
 
 import torch
@@ -42,6 +42,7 @@ def forward_intermediates(
         aggregation: Optional[str] = "sparse",
         inter_feature_normalizer: Optional[IntermediateFeatureNormalizerBase] = None,
         norm_alpha_scheme = "post-alpha",
+        block_kwargs: Dict = None,
 ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
     """ Forward features that returns intermediates.
 
@@ -65,6 +66,8 @@ def forward_intermediates(
     reshape = output_fmt == 'NCHW'
     intermediates = []
 
+    block_kwargs = block_kwargs or dict()
+
     blocks = model.blocks
 
     take_indices, max_index = _take_indices(len(blocks), indices)
@@ -90,7 +93,7 @@ def forward_intermediates(
     take_off = 0
 
     for i, blk in enumerate(blocks):
-        x = blk(x)
+        x = blk(x, **block_kwargs)
         if aggregation == "dense":
             # Arbitrarily use the rotation matrix from the final layer in the dense group
             y, alpha = inter_feature_normalizer(x, i, rot_index=take_indices[take_off], skip=num_summary_tokens)

radio_model.py

@@ -18,6 +18,7 @@ from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
 from . import eradio_model
 from .enable_spectral_reparam import configure_spectral_reparam_from_args
 from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
+from . import dual_hybrid_vit
 
 
 class Resolution(NamedTuple):
@@ -69,7 +70,7 @@ class RADIOModel(nn.Module):
         patch_gen = getattr(self.model, "patch_generator", None)
         if patch_gen is not None:
             return patch_gen.num_skip
-        elif self.model.global_pool == 'avg':
+        elif getattr(self.model, 'global_pool', None) == 'avg':
             return 0
         return 1
 
@@ -81,7 +82,7 @@ class RADIOModel(nn.Module):
         patch_gen = getattr(self.model, 'patch_generator', None)
         if patch_gen is not None:
             return patch_gen.num_cls_tokens
-        elif self.model.global_pool == 'avg':
+        elif getattr(self.model, 'global_pool', None) == 'avg':
             return 0
         return 1
 
@@ -218,7 +219,10 @@ class RADIOModel(nn.Module):
             ret = dict(backbone=ret)
             for name, adaptor in self.adaptors.items():
                 if all_summary.ndim == 3:
-                    summary = all_summary[:, adaptor.head_idx]
+                    if all_summary.shape[1] == 1:
+                        summary = all_summary[:, 0]
+                    else:
+                        summary = all_summary[:, adaptor.head_idx]
                 else:
                     summary = all_summary
                 ada_input = AdaptorInput(images=x, summary=summary.float(), features=all_feat, feature_fmt=feature_fmt, patch_size=self.patch_size)
@@ -326,10 +330,6 @@ def create_model_from_args(args) -> nn.Module:
 
     model.head = nn.Identity()
 
-    assert (
-        not args.cls_token_per_teacher or args.cpe_max_size is not None
-    ), "CPE must be enabled for multiple CLS tokens!"
-
     if args.cpe_max_size is not None:
         uq_teachers = set(t['name'] for t in args.teachers)
         enable_cpe(

vit_patch_generator.py

@@ -106,6 +106,10 @@ class ViTPatchGenerator(nn.Module):
     def num_cls_tokens(self):
         return self.cls_token.num_tokens
 
+    @property
+    def num_cls_patches(self):
+        return self.cls_token.num_patches
+
     @property
     def num_registers(self):
         return self.cls_token.num_registers
@@ -119,10 +123,6 @@ class ViTPatchGenerator(nn.Module):
             'pos_embed',
         ]
 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
-        if self.abs_pos:
-            self._load_embed(state_dict[f'{prefix}pos_embed'], self.pos_embed)
-
     def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
         if src_embed.shape != targ_embed.shape:
             src_size = int(math.sqrt(src_embed.shape[1]))
@@ -285,18 +285,3 @@ class ViTPatchLinear(nn.Linear):
             **factory
         )
         self.patch_size = patch_size
-
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
-        if self.bias is not None:
-            self.bias.data.copy_(state_dict[f'{prefix}bias'])
-
-        chk_weight = state_dict[f'{prefix}weight']
-        if chk_weight.shape != self.weight.shape:
-            src_patch_size = int(math.sqrt(chk_weight.shape[1] // 3))
-
-            assert (src_patch_size ** 2) * 3 == chk_weight.shape[1], 'Unable to interpolate non-square patch size'
-
-            chk_weight = rearrange(chk_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
-            chk_weight = F.interpolate(chk_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
-            chk_weight = rearrange(chk_weight, 'b c h w -> b (c h w)')
-        self.weight.data.copy_(chk_weight)