nvidia/C-RADIO

1from collections import defaultdict2from contextlib import contextmanager3from logging import getLogger4import math5import sys6from typing import List, Union, Iterable7 8import numpy as np9import torch10from torch import nn11 12from timm.models import VisionTransformer13from einops import rearrange14 15DEFAULT_NUM_WINDOWED = 516 17 18class VitDetArgs:19    def __init__(self,20                 window_size: int,21                 num_summary_tokens: int,22                 num_windowed: int = DEFAULT_NUM_WINDOWED,23    ):24        self.window_size = window_size25        self.num_summary_tokens = num_summary_tokens26        self.num_windowed = num_windowed27 28 29def apply_vitdet_arch(model: VisionTransformer, args: VitDetArgs):30    if isinstance(model, VisionTransformer):31        patch_embed = getattr(model, 'patch_generator', model.patch_embed)32 33        return ViTDetHook(patch_embed, model.blocks, args)34    else:35        print(f'Warning: Unable to apply VitDet aug!', file=sys.stderr)36 37 38class ViTDetHook:39    def __init__(self,40                 embedder: nn.Module,41                 blocks: nn.Sequential,42                 args: VitDetArgs,43    ):44        self.blocks = blocks45        self.num_summary_tokens = args.num_summary_tokens46        self.window_size = args.window_size47 48        self._input_resolution = None49        self._num_windows = None50        self._cls_patch = None51        self._order_cache = dict()52 53        embedder.register_forward_pre_hook(self._enter_model)54 55        # This will decide if we window-fy the patches56        # and enable vit-det for this iteration, and if so,57        # rearrange the patches for efficient mode switching58        blocks.register_forward_pre_hook(self._enter_blocks)59 60        is_global = True61        period = args.num_windowed + 162        for i, layer in enumerate(blocks[:-1]):63            ctr = i % period64            if ctr == 0:65                layer.register_forward_pre_hook(self._to_windows)66                is_global = False67            elif ctr == args.num_windowed:68                layer.register_forward_pre_hook(self._to_global)69                is_global = True70 71        # Always ensure the final layer is a global layer72        if not is_global:73            blocks[-1].register_forward_pre_hook(self._to_global)74 75        blocks.register_forward_hook(self._exit_model)76 77    def _enter_model(self, _, input: List[torch.Tensor]):78        self._input_resolution = input[0].shape[-2:]79 80    def _enter_blocks(self, _, input: List[torch.Tensor]):81        # print(f'{get_rank()} - ViTDet Window Size: {self._window_size}', file=sys.stderr)82 83        patches = input[0]84        patches = self._rearrange_patches(patches)85 86        return (patches,) + input[1:]87 88    def _to_windows(self, _, input: List[torch.Tensor]):89        patches = input[0]90 91        if self.num_summary_tokens:92            self._cls_patch = patches[:, :self.num_summary_tokens]93            patches = patches[:, self.num_summary_tokens:]94 95        patches = rearrange(96            patches, 'b (p t) c -> (b p) t c',97            p=self._num_windows, t=self.window_size ** 2,98        )99 100        return (patches,) + input[1:]101 102    def _to_global(self, _, input: List[torch.Tensor]):103        patches = input[0]104 105        patches = rearrange(106            patches, '(b p) t c -> b (p t) c',107            p=self._num_windows, t=self.window_size ** 2,108            b=patches.shape[0] // self._num_windows,109        )110 111        if self.num_summary_tokens:112            patches = torch.cat([113                self._cls_patch,114                patches,115            ], dim=1)116 117        return (patches,) + input[1:]118 119    def _exit_model(self, _, inputs: List[torch.Tensor], patches: torch.Tensor):120        # Return patches to their original order121        patch_order = self._order_cache[self._input_resolution][0]122        patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)123 124        ret_patches = torch.empty_like(patches)125        ret_patches = torch.scatter(126            ret_patches,127            dim=1,128            index=patch_order,129            src=patches,130        )131 132        return ret_patches133 134    def _rearrange_patches(self, patches: torch.Tensor):135        # We rearrange the patches so that we can efficiently136        # switch between windowed and global mode by just137        # reshaping the tensor138 139        patch_order, self._num_windows = self._order_cache.get(self._input_resolution, (None, None))140        if patch_order is None:141            num_feat_patches = patches.shape[1] - self.num_summary_tokens142            num_pixels = self._input_resolution[0] * self._input_resolution[1]143 144            patch_size = int(round(math.sqrt(num_pixels / num_feat_patches)))145            rows = self._input_resolution[-2] // patch_size146            cols = self._input_resolution[-1] // patch_size147 148            w_rows = rows // self.window_size149            w_cols = cols // self.window_size150 151            patch_order = torch.arange(0, num_feat_patches, device=patches.device)152 153            patch_order = rearrange(154                patch_order, '(wy py wx px) -> (wy wx py px)',155                wy=w_rows, wx=w_cols,156                py=self.window_size, px=self.window_size,157            )158 159            if self.num_summary_tokens:160                patch_order = torch.cat([161                    torch.arange(self.num_summary_tokens, dtype=patch_order.dtype, device=patch_order.device),162                    patch_order + self.num_summary_tokens,163                ])164 165            self._num_windows = w_rows * w_cols166            self._order_cache[self._input_resolution] = (167                patch_order,168                self._num_windows,169            )170 171        patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)172        patches = torch.gather(patches, dim=1, index=patch_order)173        return patches174