nvidia/C-RADIO

1# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from collections import namedtuple15from typing import Optional, List, Union16 17from timm.models import VisionTransformer18import torch19from transformers import PretrainedConfig, PreTrainedModel20 21 22from .common import RESOURCE_MAP, DEFAULT_VERSION23 24# Import all required modules.25from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput26from .adaptor_generic import GenericAdaptor, AdaptorBase27from .adaptor_mlp import create_mlp_from_state28from .adaptor_registry import adaptor_registry29from .cls_token import ClsToken30from .enable_cpe_support import enable_cpe31from .enable_spectral_reparam import configure_spectral_reparam_from_args32from .eradio_model import eradio33from .radio_model import create_model_from_args34from .radio_model import RADIOModel as RADIOModelBase, Resolution35from .input_conditioner import get_default_conditioner, InputConditioner36from .open_clip_adaptor import OpenCLIP_RADIO37from .vit_patch_generator import ViTPatchGenerator38from .vitdet import apply_vitdet_arch, VitDetArgs39 40# Register extra models41from .extra_timm_models import *42 43 44class RADIOConfig(PretrainedConfig):45    """Pretrained Hugging Face configuration for RADIO models."""46 47    def __init__(48        self,49        args: Optional[dict] = None,50        version: Optional[str] = DEFAULT_VERSION,51        patch_size: Optional[int] = None,52        max_resolution: Optional[int] = None,53        preferred_resolution: Optional[Resolution] = None,54        adaptor_names: Union[str, List[str]] = None,55        vitdet_window_size: Optional[int] = None,56        **kwargs,57    ):58        self.args = args59        for field in ["dtype", "amp_dtype"]:60            if self.args is not None and field in self.args:61                # Convert to a string in order to make it serializable.62                # For example for torch.float32 we will store "float32",63                # for "bfloat16" we will store "bfloat16".64                self.args[field] = str(args[field]).split(".")[-1]65        self.version = version66        resource = RESOURCE_MAP[version]67        self.patch_size = patch_size or resource.patch_size68        self.max_resolution = max_resolution or resource.max_resolution69        self.preferred_resolution = (70            preferred_resolution or resource.preferred_resolution71        )72        self.adaptor_names = adaptor_names73        self.vitdet_window_size = vitdet_window_size74        super().__init__(**kwargs)75 76 77class RADIOModel(PreTrainedModel):78    """Pretrained Hugging Face model for RADIO.79 80    This class inherits from PreTrainedModel, which provides81    HuggingFace's functionality for loading and saving models.82    """83 84    config_class = RADIOConfig85 86    def __init__(self, config):87        super().__init__(config)88 89        RADIOArgs = namedtuple("RADIOArgs", config.args.keys())90        args = RADIOArgs(**config.args)91        self.config = config92 93        model = create_model_from_args(args)94        input_conditioner: InputConditioner = get_default_conditioner()95 96        dtype = getattr(args, "dtype", torch.float32)97        if isinstance(dtype, str):98            # Convert the dtype's string representation back to a dtype.99            dtype = getattr(torch, dtype)100        model.to(dtype=dtype)101        input_conditioner.dtype = dtype102 103        summary_idxs = torch.tensor(104            [i for i, t in enumerate(args.teachers) if t.get("use_summary", True)],105            dtype=torch.int64,106        )107 108        adaptor_names = config.adaptor_names109        if adaptor_names is not None:110            raise NotImplementedError(111                f"Adaptors are not yet supported in Hugging Face models. Adaptor names: {adaptor_names}"112            )113 114        adaptors = dict()115 116        self.radio_model = RADIOModelBase(117            model,118            input_conditioner,119            summary_idxs=summary_idxs,120            patch_size=config.patch_size,121            max_resolution=config.max_resolution,122            window_size=config.vitdet_window_size,123            preferred_resolution=config.preferred_resolution,124            adaptors=adaptors,125        )126 127    @property128    def model(self) -> VisionTransformer:129        return self.radio_model.model130 131    @property132    def input_conditioner(self) -> InputConditioner:133        return self.radio_model.input_conditioner134 135    def forward(self, x: torch.Tensor):136        return self.radio_model.forward(x)137