nvidia/C-RADIOv4-1D-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 9from logging import getLogger10import math11import os12from typing import Dict, List, Optional, Union, Tuple13from types import MethodType14 15import torch16from torch import nn17from torch.nn import functional as F18from torch.nn.utils import parametrize19from torch.nn.utils.parametrizations import _SpectralNorm20 21from timm.models.vision_transformer import Attention, Mlp22 23_EPS = 1e-524 25 26class _SNReweight(_SpectralNorm):27    def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, alpha: float = 0.05, version: int = 2, **kwargs):28        super().__init__(weight, *args, **kwargs)29 30        self.alpha = alpha31        self.version = version32        self.register_buffer('_sn_version', torch.tensor(version))33 34        if init_norm_to_current:35            # This will set the numerator to match the denominator, which should preserve the original values36            init_scale = self._get_sigma(weight, n_power_iterations=20).item()37        else:38            init_scale = 1.039 40        if version == 1:41            init_value = init_scale42        elif version == 2:43            t = init_scale - alpha44            if t < _EPS:45                getLogger("spectral_reparam").warn(f'The initialized spectral norm {init_scale} is too small to be represented. Setting to {_EPS} instead.')46                t = _EPS47 48            init_value = math.log(math.exp(t) - 1)49        else:50            raise ValueError(f'Unsupported version: {version}')51 52        # Make 2D so that weight decay gets applied53        self.scale = nn.Parameter(torch.tensor([[init_value]], dtype=torch.float32, device=weight.device))54 55    # Re-implementing this because we need to make division by sigma safe56    def _get_sigma(self, weight: torch.Tensor, n_power_iterations: int = None) -> torch.Tensor:57        if not n_power_iterations:58            n_power_iterations = self.n_power_iterations59        if weight.ndim == 1:60            # Faster and more exact path, no need to approximate anything61            sigma = weight.norm()62        else:63            weight_mat = self._reshape_weight_to_matrix(weight)64            if self.training:65                self._power_method(weight_mat, n_power_iterations)66            # See above on why we need to clone67            u = self._u.clone(memory_format=torch.contiguous_format)68            v = self._v.clone(memory_format=torch.contiguous_format)69            # The proper way of computing this should be through F.bilinear, but70            # it seems to have some efficiency issues:71            # https://github.com/pytorch/pytorch/issues/5809372            sigma = torch.dot(u, torch.mv(weight_mat, v))73 74        return sigma + self.eps75 76    def forward(self, weight: torch.Tensor, *args, **kwargs):77        dtype = weight.dtype78        sigma = self._get_sigma(weight, *args, **kwargs)79 80        if self.version == 1:81            scale = self.scale82        elif self.version == 2:83            scale = F.softplus(self.scale) + self.alpha84        else:85            raise ValueError(f'Unsupported version: {self.version}')86 87        scale = scale.float() / sigma.float()88 89        y = weight * scale90 91        if dtype in (torch.float16, torch.bfloat16):92            y = y.to(dtype)93        return y94 95    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):96        version_key = f'{prefix}_sn_version'97        if version_key not in state_dict:98            self.version = 199            state_dict[version_key] = torch.tensor(1)100        return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)101 102 103class _ChunkedSNReweight(nn.Module):104    def __init__(self, weight: torch.Tensor, num_chunks: int, *args, init_norm_to_current: bool = False, **kwargs):105        super().__init__()106 107        self.num_chunks = num_chunks108        parts = weight.split(weight.shape[0] // num_chunks, dim=0)109 110        self.parts = nn.ModuleList([111            _SNReweight(p, *args, init_norm_to_current=init_norm_to_current, **kwargs)112            for p in parts113        ])114 115    def forward(self, weight: torch.Tensor, *args, **kwargs):116        parts = weight.split(weight.shape[0] // self.num_chunks, dim=0)117 118        parts = [119            fn(p)120            for fn, p in zip(self.parts, parts)121        ]122 123        return torch.cat(parts, dim=0)124 125 126class _AttnSNReweight(_ChunkedSNReweight):127    def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, renorm_values: bool = False, **kwargs):128        super().__init__(weight, 3, *args, init_norm_to_current=init_norm_to_current, **kwargs)129 130        if not renorm_values:131            self.parts[2] = nn.Identity()132 133 134def enable_spectral_reparam(model: Union[nn.Module, List[nn.Module]],135                            n_power_iterations: int = 1,136                            eps: float = 1e-6,137                            init_norm_to_current: bool = False,138                            renorm_values: bool = True,139                            renorm_mlp: bool = True,140                            state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):141    if isinstance(model, (list, tuple)):142        for i, sub in enumerate(model):143            sub_sd = state_dict_guidance[i] if isinstance(state_dict_guidance, (list, tuple)) else state_dict_guidance144            enable_spectral_reparam(sub, n_power_iterations=n_power_iterations, eps=eps,145                                    init_norm_to_current=init_norm_to_current, renorm_values=renorm_values,146                                    renorm_mlp=renorm_mlp, state_dict_guidance=sub_sd)147        return148 149    print('Enabling spectral reparametrization')150    args = dict(n_power_iterations=n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current)151    visited_prefixes = set()152 153    def is_guidance_parametrized(name: str):154        if state_dict_guidance is None:155            return True156 157        p_name = f'{name}.parametrizations'158        is_prm = any(k for k in state_dict_guidance if k.startswith(p_name) and k.endswith('_sn_version'))159        return is_prm160 161    def parametrize_linear(linear: nn.Linear):162        parametrize.register_parametrization(163            linear,164            'weight',165            _SNReweight(linear.weight, **args)166        )167 168    for name, mod in model.named_modules():169        pref = '.'.join(name.split('.')[:-1])170        if pref in visited_prefixes:171            continue172 173        if isinstance(mod, Attention) or name.endswith('.attn'):174            if is_guidance_parametrized(f'{name}.qkv'):175                parametrize.register_parametrization(176                    mod.qkv,177                    'weight',178                    _AttnSNReweight(mod.qkv.weight, renorm_values=renorm_values, **args),179                )180            if hasattr(mod, 'proj') and is_guidance_parametrized(f'{name}.proj'):181                parametrize_linear(mod.proj)182            visited_prefixes.add(name)183        elif name.endswith('mlp') and renorm_mlp and hasattr(mod, 'w12'):184            if is_guidance_parametrized(f'{name}.w12'):185                parametrize.register_parametrization(186                    mod.w12,187                    'weight',188                    _ChunkedSNReweight(mod.w12.weight, num_chunks=2, **args),189                )190            if is_guidance_parametrized(f'{name}.w3'):191                parametrize_linear(mod.w3)192            visited_prefixes.add(name)193        elif isinstance(mod, nn.Linear) and 'patch_generator' not in name and is_guidance_parametrized(name):194            parametrize_linear(mod)195 196 197def configure_spectral_reparam_from_args(model: nn.Module, args, state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):198    spectral_reparam = getattr(args, 'spectral_reparam', False)199    if isinstance(spectral_reparam, bool) and spectral_reparam:200        enable_spectral_reparam(model, init_norm_to_current=True, state_dict_guidance=state_dict_guidance)201    elif isinstance(spectral_reparam, dict):202        enable_spectral_reparam(203            model,204            n_power_iterations=spectral_reparam.get('n_power_iterations', 1),205            eps=spectral_reparam.get('eps', 1e-12),206            init_norm_to_current=True,207            state_dict_guidance=state_dict_guidance,208        )209 210 211def disable_spectral_reparam(model: nn.Module):212    print('Disabling spectral reparametrization')213    for name, mod in model.named_modules():214        if parametrize.is_parametrized(mod):215            parametrize.remove_parametrizations(mod, 'weight')216            pass217 218 219 220if __name__ == '__main__':221    import argparse222    from . import radio_model as create_model223 224    parser = argparse.ArgumentParser(description='Remove parametrization from state dict')225    parser.add_argument('--checkpoint', type=str, required=True, help='The checkpoint to load')226    parser.add_argument('--output', type=str, default='', help='Where to store the checkpoint')227    parser.add_argument('--release', default=False, action='store_true', help='Prune extraneous checkpoint fields')228    parser.add_argument('--strict', default=False, action='store_true', help='Strictly load the state dict')229 230    args = parser.parse_args()231 232    if not args.output:233        chk_dir, chk_name = os.path.split(args.checkpoint)234        args.output = os.path.join(chk_dir, f'clean_{chk_name}')235        print(f'Set output to "{args.output}"')236 237    chk = torch.load(args.checkpoint, map_location='cpu', mmap=True)238 239    model = create_model.create_model_from_args(chk['args'])240 241    key = 'base_model.'242    mod_state = dict()243    extra_state = dict()244    for k, v in chk['state_dict'].items():245        if k.startswith(key):246            mod_state[k[len(key):]] = v247        else:248            extra_state[k] = v249 250    chk_load_info = model.load_state_dict(mod_state, strict=args.strict)251    if chk_load_info.unexpected_keys or chk_load_info.missing_keys:252        print(chk_load_info)253 254    if chk['args'].spectral_reparam:255        disable_spectral_reparam(model)256 257    if hasattr(chk['args'], 'dtype'):258        model.to(dtype=chk['args'].dtype)259 260    mod_state = model.state_dict()261    final_state = dict()262    final_state.update({f'{key}{k}': v for k, v in mod_state.items()})263    final_state.update(extra_state)264 265    chk['state_dict'] = final_state266    chk['args'].spectral_reparam = False267 268    if args.release:269        chk = {270            'arch': chk['arch'],271            'epoch': chk['epoch'],272            'state_dict': chk['state_dict'],273            'args': chk['args'],274        }275 276    torch.save(chk, args.output)277    pass278