nvidia/C-RADIO

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 9from torch import nn10 11from timm.models import register_model12from timm.models.vision_transformer import VisionTransformer, _create_vision_transformer, Mlp13 14 15@register_model16def vit_tiny_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:17    """ ViT-Tiny (Vit-Ti/16)18    """19    model_args = dict(patch_size=14, embed_dim=192, depth=12, num_heads=3)20    model = _create_vision_transformer('vit_tiny_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))21    return model22 23 24@register_model25def vit_small_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:26    """ ViT-Small (ViT-S/16)27    """28    model_args = dict(patch_size=14, embed_dim=384, depth=12, num_heads=6)29    model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))30    return model31 32 33@register_model34def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:35    """ ViT-Base (ViT-B/14) from original paper (https://arxiv.org/abs/2010.11929).36    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.37    """38    model_args = dict(patch_size=14, embed_dim=768, depth=12, num_heads=12)39    model = _create_vision_transformer('vit_base_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))40    return model41 42 43@register_model44def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:45    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).46    """47    model_args = dict(patch_size=16, embed_dim=1280, depth=32, num_heads=16)48    if pretrained:49        # There is no pretrained version of ViT-H/16, but we can adapt a ViT-H/14 for this purpose50        model = _create_vision_transformer('vit_huge_patch14_clip_336', pretrained=True, **dict(model_args, pre_norm=True, **kwargs))51    else:52        model = _create_vision_transformer('vit_huge_patch16_224', pretrained=False, **dict(model_args, **kwargs))53    return model54 55 56@register_model57def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransformer:58    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).59    """60    model = vit_huge_patch16_224(pretrained=pretrained, **kwargs)61 62    for m in model.modules():63        if isinstance(m, Mlp) and not isinstance(m.norm, nn.LayerNorm):64            m.norm = nn.LayerNorm(m.fc1.out_features)65 66    return model67