nvidia/C-RADIOv2-VLM-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, 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) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:46    """ Forward features that returns intermediates.47 48    The Dense layer aggregation method is inspired from the paper: "Dense Connector for MLLMs"49    by Yao, Huanjin et al. (2024). arXiv preprint arXiv:2405.13800}50 51    Args:52        x: Input image tensor53        indices: Take last n blocks if int, select matching indices if sequence54        return_prefix_tokens: Return both prefix and spatial intermediate tokens55        norm: Apply norm layer to all intermediates56        stop_early: Stop iterating over blocks when last desired intermediate hit57        output_fmt: Shape of intermediate feature outputs58        intermediates_only: Only return intermediate features59        aggregation: intermediate layer aggregation method (sparse or dense)60        norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha")61    Returns:62    """63    assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'64    assert aggregation in ('sparse', 'dense'), 'Aggregation must be one of sparse or dense.'65    reshape = output_fmt == 'NCHW'66    intermediates = []67 68    blocks = model.blocks69 70    take_indices, max_index = _take_indices(len(blocks), indices)71    take_indices = sorted(take_indices)72    # forward pass73    B, _, height, width = x.shape74 75    x = patch_extractor(x)76 77    if stop_early:78        blocks = blocks[:max_index + 1]79 80    if inter_feature_normalizer is None or norm_alpha_scheme == 'none':81        inter_feature_normalizer = NullIntermediateFeatureNormalizer.get_instance(x.dtype, x.device)82 83    assert norm_alpha_scheme in ('none', 'pre-alpha', 'post-alpha'), f'Unsupported alpha scheme: {norm_alpha_scheme}'84    post_alpha_scheme = norm_alpha_scheme == 'post-alpha'85 86    accumulator = 087    alpha_sum = 088    num_accumulated = 089 90    take_off = 091 92    for i, blk in enumerate(blocks):93        x = blk(x)94        if aggregation == "dense":95            # Arbitrarily use the rotation matrix from the final layer in the dense group96            y, alpha = inter_feature_normalizer(x, i, rot_index=take_indices[take_off], skip=num_summary_tokens)97            if post_alpha_scheme:98                accumulator = accumulator + y99                alpha_sum = alpha_sum + alpha100            else:101                accumulator = accumulator + (alpha * y)102                alpha_sum += 1103            num_accumulated += 1104        if i == take_indices[take_off]:105            if aggregation == "dense":106                alpha = alpha_sum / num_accumulated107                x_ = alpha * accumulator / num_accumulated108                num_accumulated = 0109                accumulator = 0110                alpha_sum = 0111            else:112                 y, alpha = inter_feature_normalizer(x, i, skip=num_summary_tokens)113                 x_ = alpha * y114            # normalize intermediates with final norm layer if enabled115            intermediates.append(norm(x_))116            take_off = min(take_off + 1, len(take_indices) - 1)117 118    # process intermediates119 120    # split prefix (e.g. class, distill) and spatial feature tokens121    prefix_tokens = [y[:, :num_cls_tokens] for y in intermediates]122    intermediates = [y[:, num_summary_tokens:] for y in intermediates]123 124    if reshape:125        # reshape to BCHW output format126        H = height // model.patch_size127        W = width // model.patch_size128        intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]129    if not torch.jit.is_scripting() and return_prefix_tokens:130        # return_prefix not support in torchscript due to poor type handling131        intermediates = list(zip(prefix_tokens, intermediates))132    if intermediates_only:133        return intermediates134    x = norm(x)135    return x, intermediates136