nvidia/C-RADIOv4-H

1# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9import math10import warnings11 12import torch13from torch import nn14from torch.nn import functional as F15 16from timm.models import register_model, PretrainedCfg17from timm.models.vision_transformer import (18    VisionTransformer,19    _create_vision_transformer as _timm_create_vision_transformer,20    Mlp,21    Block,22    LayerScale as TIMMLayerScale,23)24 25# Import these to also register them26from . import dinov2_arch27 28 29@register_model30def vit_tiny_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:31    """ ViT-Tiny (Vit-Ti/16)32    """33    model_args = dict(patch_size=14, embed_dim=192, depth=12, num_heads=3)34    model = _create_vision_transformer('vit_tiny_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))35    return model36 37 38@register_model39def vit_small_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:40    """ ViT-Small (ViT-S/16)41    """42    model_args = dict(patch_size=14, embed_dim=384, depth=12, num_heads=6)43    model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))44    return model45 46 47@register_model48def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:49    """ ViT-Base (ViT-B/14) from original paper (https://arxiv.org/abs/2010.11929).50    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.51    """52    model_args = dict(patch_size=14, embed_dim=768, depth=12, num_heads=12)53    model = _create_vision_transformer('vit_base_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))54    return model55 56 57@register_model58def vit_base_patch16_v2_224(pretrained=False, **kwargs) -> VisionTransformer:59    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).60    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.61    """62    model_args = dict(63        patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5,64        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 1465    )66    model = _create_vision_transformer(67        'vit_base_patch14_reg4_dinov2', pretrained=False, **dict(model_args, **kwargs))68    return model69 70 71@register_model72def vit_large_patch16_v2_224(pretrained: bool = False, **kwargs) -> VisionTransformer:73    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).74    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.75    """76    name = 'vit_large_patch14_reg4_dinov2'77    model_args = dict(78        patch_size=16, embed_dim=1024, depth=24, num_heads=16, init_values=1e-5,79        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 1480    )81    model = _create_vision_transformer(name, pretrained=False, **dict(model_args, **kwargs))82 83    return model84 85 86@register_model87def vit_so400m_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:88    """ ViT model matching the architecture of the So400M model from89    "Scaling Vision Transformers to 400 Million Parameters" (https://arxiv.org/abs/2302.05442).90    """91    if pretrained:92        raise ValueError('There is no pretrained weights for vit_so400m_patch16_224')93    mlp_ratio = 4304 / 115294 95    model_args = dict(patch_size=16, embed_dim=1152, depth=27, num_heads=16, mlp_ratio=mlp_ratio)96    model = _create_vision_transformer('vit_so400m_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))97    return model98 99 100@register_model101def vit_so400m_v2_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:102    """ ViT model matching the architecture of the So400M model from103    "Scaling Vision Transformers to 400 Million Parameters" (https://arxiv.org/abs/2302.05442).104    """105    if pretrained:106        raise ValueError('There is no pretrained weights for vit_so400m_patch16_224')107 108    normal_target = 4304109    # TP4 requires channels to be a multiple of 4, and then within that, FP8 requires a multiple of 8,110    # thus, a multiple of 32 is required.111    tp4_fp8_safe_target = ((normal_target + 31) // 32) * 32112 113    mlp_ratio = tp4_fp8_safe_target / 1152114 115    model_args = dict(patch_size=16, embed_dim=1152, depth=27, num_heads=16, mlp_ratio=mlp_ratio)116    model = _create_vision_transformer('vit_so400m_v2_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))117    return model118 119 120@register_model121def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:122    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).123    """124    model_args = dict(patch_size=16, embed_dim=1280, depth=32, num_heads=16)125    if pretrained:126        # There is no pretrained version of ViT-H/16, but we can adapt a ViT-H/14 for this purpose127        model = _create_vision_transformer('vit_huge_patch14_224', pretrained=True, **dict(model_args, **kwargs))128    else:129        model = _create_vision_transformer('vit_huge_patch16_224', pretrained=False, **dict(model_args, **kwargs))130    return model131 132 133@register_model134def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransformer:135    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).136    """137    model = vit_huge_patch16_224(pretrained=pretrained, **kwargs)138 139    for m in model.modules():140        if isinstance(m, Mlp) and not isinstance(m.norm, nn.LayerNorm):141            m.norm = nn.LayerNorm(m.fc1.out_features)142 143    return model144 145 146@register_model147def vit_giant_patch16_224(pretrained=False, scaled_ln: bool = False, **kwargs) -> VisionTransformer:148    """ ViT-giant model (ViT-g/16) from original paper (https://arxiv.org/abs/2010.11929).149    """150    model_args = dict(patch_size=16, embed_dim=1536, depth=40, num_heads=24)151    model = _create_vision_transformer('vit_giant_patch16_224', pretrained=False, **dict(model_args, **kwargs))152    if scaled_ln:153        _apply_scaled_ln(model)154    return model155 156 157@register_model158def vit_bigG_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:159    model_args = dict(patch_size=14, embed_dim=1664, depth=48, num_heads=16, init_values=1e-6)160    model = _create_vision_transformer('vit_bigG_patch14', pretrained=False, **dict(model_args, **kwargs))161    return model162 163 164def _create_vision_transformer(*args, **kwargs):165    if kwargs.get('pretrained_cfg', None) is None:166        # This prevents the warning from being emitted167        kwargs['pretrained_cfg'] = PretrainedCfg()168 169    model = _timm_create_vision_transformer(*args, **kwargs)170    _patch_layer_scale(model)171    return model172 173 174def _patch_layer_scale(model: VisionTransformer):175    def replace_ls(old_ls: TIMMLayerScale):176        new_ls = dinov2_arch.LayerScale(old_ls.gamma.shape[0], inplace=old_ls.inplace)177        new_ls.load_state_dict(old_ls.state_dict())178        return new_ls179 180    # Monkey patch: Replace TIMM's LayerScale with our modified DINOv2 one, that uses a param name181    # other than gamma, so that HFHub doesn't mess with it!182    for mod in model.modules():183        if isinstance(mod, Block):184            if isinstance(mod.ls1, TIMMLayerScale):185                mod.ls1 = replace_ls(mod.ls1)186            if isinstance(mod.ls2, TIMMLayerScale):187                mod.ls2 = replace_ls(mod.ls2)188    pass189 190 191class ScaledLayerNorm(nn.LayerNorm):192    '''193    https://arxiv.org/pdf/2502.05795v1194    '''195    def __init__(self, ln_base: nn.LayerNorm, depth: int = 0):196        super().__init__(ln_base.normalized_shape, eps=ln_base.eps, elementwise_affine=ln_base.elementwise_affine)197        self.load_state_dict(ln_base.state_dict())198        self.register_buffer('ln_scale', torch.tensor(1.0 / math.sqrt(depth)), persistent=False)199 200    def forward(self, x):201        y = super().forward(x)202        y = y * self.ln_scale203        return y204 205 206class DyT(nn.Module):207    def __init__(self, C: int, init_alpha: float):208        super().__init__()209        self.alpha = nn.Parameter(torch.full((1,), init_alpha))210        self.gamma = nn.Parameter(torch.ones(C))211        self.beta = nn.Parameter(torch.zeros(C))212 213    def forward(self, x: torch.Tensor):214        x = F.tanh(self.alpha * x)215        return self.gamma * x + self.beta216 217@register_model218def vit_large_dyt_patch16_224(pretrained: bool = False, **kwargs) -> VisionTransformer:219    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).220    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.221    """222    model_args = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16)223    model = _create_vision_transformer('vit_large_dyt_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))224 225    def _replace_ln_with_dyt(ln: nn.LayerNorm, depth: int):226        return DyT(ln.normalized_shape[0], init_alpha=0.9)227    _replace_ln(model, _replace_ln_with_dyt)228 229    return model230 231 232def _apply_scaled_ln(model: VisionTransformer):233    warnings.warn('Post-LayerNorm scaling activated!')234 235    _replace_ln(model, lambda ln, depth: ScaledLayerNorm(ln, depth=depth))236 237def _replace_ln(model: VisionTransformer, fn):238    def _inner_replace_ln(block: Block, depth: int, key: str):239        prev = getattr(block, key)240        if isinstance(prev, nn.LayerNorm):241            setattr(block, key, fn(prev, depth=depth))242 243    for i, block in enumerate(model.blocks):244        _inner_replace_ln(block, i + 1, 'norm1')245        _inner_replace_ln(block, i + 1, 'norm2')246