nvidia/C-RADIOv4-H

1# Copyright (c) 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 argparse import Namespace9 10import torch11from torch import nn12import torch.nn.functional as F13 14from .adaptor_base import AdaptorBase, AdaptorInput, RadioOutput15from .adaptor_module_factory import create_mlp_from_state, create_mlp_from_config16 17 18class GenericAdaptor(AdaptorBase):19    def __init__(self, main_config: Namespace, adaptor_config, state, mlp_config=None):20        super().__init__()21 22        summary_mlp_version = main_config.mlp_version23        feature_mlp_version = getattr(main_config, 'spatial_mlp_version', None) or summary_mlp_version24 25        extra_args = dict()26        ups = None27        ups_rank = None28        if adaptor_config is not None:29            ups = adaptor_config.get('fd_upsample_factor', None)30            ups_rank = adaptor_config.get('fd_upsample_rank', None)31            summary_mlp_version = adaptor_config.get('mlp_version', summary_mlp_version)32            feature_mlp_version = adaptor_config.get('spatial_mlp_version', feature_mlp_version)33        elif mlp_config is not None:34            ups = mlp_config["feature"].get('upsample_factor', None)35            ups_rank = mlp_config["feature"].get('upsample_rank', None)36        if ups is not None:37            extra_args['upsample_factor'] = ups38            extra_args['upsample_rank'] = ups_rank39 40        if state is not None:41            spectral_heads = getattr(main_config, 'spectral_heads', False)42            self.head_mlp = create_mlp_from_state(summary_mlp_version, state, 'summary.', spectral_weights=spectral_heads, is_summary=True)43            self.feat_mlp = create_mlp_from_state(feature_mlp_version, state, 'feature.', spectral_weights=spectral_heads, is_summary=False, **extra_args)44        else:45            assert mlp_config is not None, "Config must not be None if state is None"46 47            self.head_mlp =  create_mlp_from_config(48                summary_mlp_version,49                mlp_config["summary"]["input_dim"],50                mlp_config["summary"]["hidden_dim"],51                mlp_config["summary"]["output_dim"],52                mlp_config["summary"]["num_inner"],53                is_summary=True,54            )55            self.feat_mlp = create_mlp_from_config(56                feature_mlp_version,57                mlp_config["feature"]["input_dim"],58                mlp_config["feature"]["hidden_dim"],59                mlp_config["feature"]["output_dim"],60                mlp_config["feature"]["num_inner"],61                is_summary=False,62                **extra_args63            )64 65    def forward(self, input: AdaptorInput) -> RadioOutput:66        # Convert input'd type to the type of the first parameter of the adaptor.67        first_param = next(self.parameters())68        summary = self.head_mlp(input.summary.to(dtype=first_param.dtype)).to(dtype=input.summary.dtype)69        feat = self.feat_mlp(input.features.to(dtype=first_param.dtype), images=input.images, patch_size=input.patch_size).to(dtype=input.features.dtype)70 71        if input.feature_fmt == 'NCHW':72            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])73                        .permute(0, 3, 1, 2)74            )75 76        return RadioOutput(summary, feat)77