nvidia/C-RADIOv4-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.8from collections import namedtuple9from typing import NamedTuple, Optional, Tuple10import torch11from torch import nn12 13 14def _run_kernel(x: torch.Tensor, mean: torch.Tensor, tx: torch.Tensor):15    if x.ndim <= 3:16        x = x - mean17        x = x @ tx.T18    elif x.ndim == 4:19        x = x - mean.reshape(1, -1, 1, 1)20        kernel = tx.reshape(*tx.shape, 1, 1)21        x = torch.nn.functional.conv2d(x, weight=kernel, bias=None, stride=1, padding=0)22    else:23        raise ValueError(f'Unsupported input dimension: {x.ndim}, shape: {x.shape}')24    return x25 26 27class FeatureNormalizer(nn.Module):28    def __init__(self, embed_dim: int, dtype: torch.dtype = torch.float32):29        super().__init__()30 31        self.register_buffer('mean', torch.zeros(embed_dim, dtype=dtype))32        self.register_buffer('tx', torch.eye(embed_dim, dtype=dtype))33 34    def forward(self, x: torch.Tensor) -> torch.Tensor:35        x = _run_kernel(x, self.mean, self.tx)36        return x37 38 39class InterFeatState(NamedTuple):40    y: torch.Tensor41    alpha: torch.Tensor42 43 44class IntermediateFeatureNormalizerBase(nn.Module):45    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:46        raise NotImplementedError()47 48 49class IntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):50    def __init__(self, num_intermediates: int, embed_dim: int, rot_per_layer: bool = False, dtype: torch.dtype = torch.float32):51        super().__init__()52        self.register_buffer('alphas', torch.ones(num_intermediates, dtype=dtype))53 54        rot = torch.eye(embed_dim, dtype=dtype)55        if rot_per_layer:56            rot = rot.unsqueeze(0).repeat(num_intermediates, 1, 1)57 58        self.register_buffer('rotation', rot.contiguous())59        self.register_buffer('means', torch.zeros(num_intermediates, embed_dim, dtype=dtype))60 61    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:62        if rot_index is None:63            rot_index = index64 65        if skip:66            assert x.ndim == 3, f'Cannot use the `skip` parameter when the `x` tensor isn\'t 3-dimensional.'67            prefix, x = x[:, :skip], x[:, skip:]68 69        rotation = self._get_rotation(rot_index)70        y = _run_kernel(x, self.means[index], rotation)71 72        alpha = self.alphas[index]73        if skip:74            alpha = torch.cat([75                torch.ones(skip, dtype=alpha.dtype, device=alpha.device),76                alpha[None].expand(y.shape[1]),77            ]).reshape(1, -1, 1)78            y = torch.cat([prefix, y], dim=1)79        else:80            if x.ndim == 3:81                alpha = alpha.reshape(1, 1, 1).expand(1, y.shape[1], 1)82            elif x.ndim == 4:83                alpha = alpha.reshape(1, 1, 1, 1).expand(1, 1, *y.shape[2:])84            else:85                raise ValueError(f'Unsupported input dimension: {x.ndim}')86 87        return InterFeatState(y, alpha)88 89    def _get_rotation(self, rot_index: int) -> torch.Tensor:90        if self.rotation.ndim == 2:91            return self.rotation92        return self.rotation[rot_index]93 94 95class NullIntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):96    instances = dict()97 98    def __init__(self, dtype: torch.dtype, device: torch.device):99        super().__init__()100        self.register_buffer('alpha', torch.tensor(1, dtype=dtype, device=device))101 102    @staticmethod103    def get_instance(dtype: torch.dtype, device: torch.device):104        instance = NullIntermediateFeatureNormalizer.instances.get((dtype, device), None)105        if instance is None:106            instance = NullIntermediateFeatureNormalizer(dtype, device)107            NullIntermediateFeatureNormalizer.instances[(dtype, device)] = instance108        return instance109 110    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:111        return InterFeatState(x, self.alpha)112