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.8 9from typing import Callable, Dict, List, Optional, Set, Tuple, Union, Any, Iterable10from types import MethodType11 12import torch13from torch import nn14 15from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer16 17 18def _take_indices(19        num_blocks: int,20        n: Optional[Union[int, List[int], Tuple[int]]],21) -> Tuple[Set[int], int]:22    if isinstance(n, int):23        assert n >= 024        take_indices = {x for x in range(num_blocks - n, num_blocks)}25    else:26        take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}27    return take_indices, max(take_indices)28 29 30def forward_intermediates(31        model: nn.Module,32        patch_extractor: Callable[[torch.Tensor], torch.Tensor],33        norm: nn.Module,34        num_summary_tokens: int,35        num_cls_tokens: int,36        x: torch.Tensor,37        indices: Optional[Union[int, List[int], Tuple[int]]] = None,38        return_prefix_tokens: bool = False,39        stop_early: bool = False,40        output_fmt: str = 'NCHW',41        intermediates_only: bool = False,42        aggregation: Optional[str] = "sparse",43        inter_feature_normalizer: Optional[IntermediateFeatureNormalizerBase] = None,44        norm_alpha_scheme = "post-alpha",45        block_kwargs: Dict = None,46) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:47    """ Forward features that returns intermediates.48 49    The Dense layer aggregation method is inspired from the paper: "Dense Connector for MLLMs"50    by Yao, Huanjin et al. (2024). arXiv preprint arXiv:2405.13800}51 52    Args:53        x: Input image tensor54        indices: Take last n blocks if int, select matching indices if sequence55        return_prefix_tokens: Return both prefix and spatial intermediate tokens56        norm: Apply norm layer to all intermediates57        stop_early: Stop iterating over blocks when last desired intermediate hit58        output_fmt: Shape of intermediate feature outputs59        intermediates_only: Only return intermediate features60        aggregation: intermediate layer aggregation method (sparse or dense)61        norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha")62    Returns:63    """64    assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'65    assert aggregation in ('sparse', 'dense'), 'Aggregation must be one of sparse or dense.'66    reshape = output_fmt == 'NCHW'67    intermediates = []68 69    block_kwargs = block_kwargs or dict()70 71    blocks = model.blocks72 73    take_indices, max_index = _take_indices(len(blocks), indices)74    take_indices = sorted(take_indices)75    # forward pass76    B, _, height, width = x.shape77 78    x = patch_extractor(x)79 80    if stop_early:81        blocks = blocks[:max_index + 1]82 83    if inter_feature_normalizer is None or norm_alpha_scheme == 'none':84        inter_feature_normalizer = NullIntermediateFeatureNormalizer.get_instance(x.dtype, x.device)85 86    assert norm_alpha_scheme in ('none', 'pre-alpha', 'post-alpha'), f'Unsupported alpha scheme: {norm_alpha_scheme}'87    post_alpha_scheme = norm_alpha_scheme == 'post-alpha'88 89    accumulator = 090    alpha_sum = 091    num_accumulated = 092 93    take_off = 094 95    for i, blk in enumerate(blocks):96        x = blk(x, **block_kwargs)97        if aggregation == "dense":98            # Arbitrarily use the rotation matrix from the final layer in the dense group99            y, alpha = inter_feature_normalizer(x, i, rot_index=take_indices[take_off], skip=num_summary_tokens)100            if post_alpha_scheme:101                accumulator = accumulator + y102                alpha_sum = alpha_sum + alpha103            else:104                accumulator = accumulator + (alpha * y)105                alpha_sum += 1106            num_accumulated += 1107        if i == take_indices[take_off]:108            if aggregation == "dense":109                alpha = alpha_sum / num_accumulated110                x_ = alpha * accumulator / num_accumulated111                num_accumulated = 0112                accumulator = 0113                alpha_sum = 0114            else:115                 y, alpha = inter_feature_normalizer(x, i, skip=num_summary_tokens)116                 x_ = alpha * y117            # normalize intermediates with final norm layer if enabled118            intermediates.append(norm(x_))119            take_off = min(take_off + 1, len(take_indices) - 1)120 121    # process intermediates122 123    # split prefix (e.g. class, distill) and spatial feature tokens124    prefix_tokens = [y[:, :num_cls_tokens] for y in intermediates]125    intermediates = [y[:, num_summary_tokens:] for y in intermediates]126 127    if reshape:128        # reshape to BCHW output format129        H = height // model.patch_size130        W = width // model.patch_size131        intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]132    if not torch.jit.is_scripting() and return_prefix_tokens:133        # return_prefix not support in torchscript due to poor type handling134        intermediates = list(zip(prefix_tokens, intermediates))135    if intermediates_only:136        return intermediates137    x = norm(x)138    return x, intermediates139