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.8from typing import Optional, Callable, Union, Tuple, Any, Dict, NamedTuple9 10import torch11from torch import nn12 13from timm.models import create_model, VisionTransformer14 15from .enable_cpe_support import enable_cpe16from .input_conditioner import InputConditioner17# Register extra models18from . import extra_timm_models19from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput20from . import eradio_model21from .enable_spectral_reparam import configure_spectral_reparam_from_args22 23 24class Resolution(NamedTuple):25    height: int26    width: int27 28 29class RADIOModel(nn.Module):30    def __init__(31        self,32        model: nn.Module,33        input_conditioner: InputConditioner,34        patch_size: int,35        max_resolution: int,36        preferred_resolution: Resolution,37        summary_idxs: Optional[torch.Tensor] = None,38        window_size: int = None,39        adaptors: Dict[str, AdaptorBase] = None,40    ):41        super().__init__()42 43        self.model = model44        self.input_conditioner = input_conditioner45        if summary_idxs is not None:46            self.register_buffer('summary_idxs', summary_idxs)47        else:48            self.summary_idxs = None49 50        self._preferred_resolution = preferred_resolution51        self._patch_size = patch_size52        self._max_resolution = max_resolution53        self._window_size = window_size54 55        adaptors = adaptors or dict()56        self.adaptors = nn.ModuleDict(adaptors)57 58    @property59    def num_summary_tokens(self) -> int:60        patch_gen = getattr(self.model, "patch_generator", None)61        if patch_gen is not None:62            return patch_gen.num_skip63        elif self.model.global_pool == 'avg':64            return 065        return 166 67    @property68    def patch_size(self) -> int:69        return self._patch_size70 71    @property72    def max_resolution(self) -> int:73        return self._max_resolution74 75    @property76    def preferred_resolution(self) -> Resolution:77        return self._preferred_resolution78 79    @property80    def window_size(self) -> int:81        return self._window_size82 83    @property84    def min_resolution_step(self) -> int:85        res = self.patch_size86        if self.window_size is not None:87            res *= self.window_size88        return res89 90    def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:91        ret = self.input_conditioner92        self.input_conditioner = nn.Identity()93        return ret94 95    def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:96        height = int(round(height / self.min_resolution_step) * self.min_resolution_step)97        width = int(round(width / self.min_resolution_step) * self.min_resolution_step)98 99        height = max(height, self.min_resolution_step)100        width = max(width, self.min_resolution_step)101 102        return Resolution(height=height, width=width)103 104    def switch_to_deploy(self):105        fn = getattr(self.model, 'switch_to_deploy', None)106        if fn is not None:107            fn()108 109    def forward(self, x: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:110        x = self.input_conditioner(x)111        y = self.model.forward_features(x)112 113        if isinstance(self.model, VisionTransformer):114            patch_gen = getattr(self.model, "patch_generator", None)115            if patch_gen is not None:116                all_summary = y[:, : patch_gen.num_cls_tokens]117                if self.summary_idxs is not None:118                    bb_summary = all_summary[:, self.summary_idxs]119                else:120                    bb_summary = all_summary121                all_feat = y[:, patch_gen.num_skip :]122            elif self.model.global_pool == "avg":123                all_summary = y[:, self.model.num_prefix_tokens :].mean(dim=1)124                bb_summary = all_summary125                all_feat = y126            else:127                all_summary = y[:, 0]128                bb_summary = all_summary129                all_feat = y[:, 1:]130        elif isinstance(self.model, eradio_model.FasterViT):131            _, f = y132            all_feat = f.flatten(2).transpose(1, 2)133            all_summary = all_feat.mean(dim=1)134            bb_summary = all_summary135        elif isinstance(y, (list, tuple)):136            all_summary, all_feat = y137            bb_summary = all_summary138        else:139            raise ValueError("Unsupported model type")140 141        all_feat = all_feat.float()142        ret = RadioOutput(bb_summary.flatten(1), all_feat).to(torch.float32)143        if self.adaptors:144            ret = dict(backbone=ret)145            for name, adaptor in self.adaptors.items():146                if all_summary.ndim == 3:147                    summary = all_summary[:, adaptor.head_idx]148                else:149                    summary = all_summary150                ada_input = AdaptorInput(images=x, summary=summary.float(), features=all_feat)151                v = adaptor(ada_input).to(torch.float32)152                ret[name] = v153 154        return ret155 156 157def create_model_from_args(args) -> nn.Module:158    in_chans = 3159    if args.in_chans is not None:160        in_chans = args.in_chans161    elif args.input_size is not None:162        in_chans = args.input_size[0]163 164    # Skip weight initialization unless it's explicitly requested.165    weight_init = args.model_kwargs.pop("weight_init", "skip")166 167    model = create_model(168        args.model,169        pretrained=args.pretrained,170        in_chans=in_chans,171        num_classes=args.num_classes,172        drop_rate=args.drop,173        drop_path_rate=args.drop_path,174        drop_block_rate=args.drop_block,175        global_pool=args.gp,176        bn_momentum=args.bn_momentum,177        bn_eps=args.bn_eps,178        scriptable=args.torchscript,179        checkpoint_path=args.initial_checkpoint,180        weight_init=weight_init,181        **args.model_kwargs,182    )183 184    if hasattr(model, 'norm') and not getattr(args, 'model_norm', False):185        model.norm = nn.Identity()186 187    model.head = nn.Identity()188 189    assert (190        not args.cls_token_per_teacher or args.cpe_max_size is not None191    ), "CPE must be enabled for multiple CLS tokens!"192 193    if args.cpe_max_size is not None:194        enable_cpe(195            model,196            args.cpe_max_size,197            num_cls_tokens=len(args.teachers) if args.cls_token_per_teacher else 1,198            register_multiple=args.register_multiple,199        )200 201    if args.spectral_reparam:202        configure_spectral_reparam_from_args(model, args)203 204    return model205