InPeerReview/RemoteSensingChangeDetection-RSCD.HA2F

1# Copyright (c) Meta Platforms, Inc. and affiliates.2#3# This source code is licensed under the Apache License, Version 2.04# found in the LICENSE file in the root directory of this source tree.5 6# References:7#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py8#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py9 10import logging11import os12from typing import Callable, List, Any, Tuple, Dict13import warnings14 15import torch16from torch import nn, Tensor17 18from .attention import Attention, MemEffAttention19from .drop_path import DropPath20from .layer_scale import LayerScale21from .mlp import Mlp22 23 24logger = logging.getLogger("dinov2")25 26 27XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None28try:29    if XFORMERS_ENABLED:30        from xformers.ops import fmha, scaled_index_add, index_select_cat31 32        XFORMERS_AVAILABLE = True33        warnings.warn("xFormers is available (Block)")34    else:35        warnings.warn("xFormers is disabled (Block)")36        raise ImportError37except ImportError:38    XFORMERS_AVAILABLE = False39 40    warnings.warn("xFormers is not available (Block)")41 42 43class Block(nn.Module):44    def __init__(45        self,46        dim: int,47        num_heads: int,48        mlp_ratio: float = 4.0,49        qkv_bias: bool = False,50        proj_bias: bool = True,51        ffn_bias: bool = True,52        drop: float = 0.0,53        attn_drop: float = 0.0,54        init_values=None,55        drop_path: float = 0.0,56        act_layer: Callable[..., nn.Module] = nn.GELU,57        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,58        attn_class: Callable[..., nn.Module] = Attention,59        ffn_layer: Callable[..., nn.Module] = Mlp,60    ) -> None:61        super().__init__()62        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")63        self.norm1 = norm_layer(dim)64        self.attn = attn_class(65            dim,66            num_heads=num_heads,67            qkv_bias=qkv_bias,68            proj_bias=proj_bias,69            attn_drop=attn_drop,70            proj_drop=drop,71        )72        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()73        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()74 75        self.norm2 = norm_layer(dim)76        mlp_hidden_dim = int(dim * mlp_ratio)77        self.mlp = ffn_layer(78            in_features=dim,79            hidden_features=mlp_hidden_dim,80            act_layer=act_layer,81            drop=drop,82            bias=ffn_bias,83        )84        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()85        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()86 87        self.sample_drop_ratio = drop_path88 89    def forward(self, x: Tensor) -> Tensor:90        def attn_residual_func(x: Tensor) -> Tensor:91            return self.ls1(self.attn(self.norm1(x)))92 93        def ffn_residual_func(x: Tensor) -> Tensor:94            return self.ls2(self.mlp(self.norm2(x)))95 96        if self.training and self.sample_drop_ratio > 0.1:97            # the overhead is compensated only for a drop path rate larger than 0.198            x = drop_add_residual_stochastic_depth(99                x,100                residual_func=attn_residual_func,101                sample_drop_ratio=self.sample_drop_ratio,102            )103            x = drop_add_residual_stochastic_depth(104                x,105                residual_func=ffn_residual_func,106                sample_drop_ratio=self.sample_drop_ratio,107            )108        elif self.training and self.sample_drop_ratio > 0.0:109            x = x + self.drop_path1(attn_residual_func(x))110            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2111        else:112            x = x + attn_residual_func(x)113            x = x + ffn_residual_func(x)114        return x115 116 117def drop_add_residual_stochastic_depth(118    x: Tensor,119    residual_func: Callable[[Tensor], Tensor],120    sample_drop_ratio: float = 0.0,121) -> Tensor:122    # 1) extract subset using permutation123    b, n, d = x.shape124    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)125    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]126    x_subset = x[brange]127 128    # 2) apply residual_func to get residual129    residual = residual_func(x_subset)130 131    x_flat = x.flatten(1)132    residual = residual.flatten(1)133 134    residual_scale_factor = b / sample_subset_size135 136    # 3) add the residual137    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)138    return x_plus_residual.view_as(x)139 140 141def get_branges_scales(x, sample_drop_ratio=0.0):142    b, n, d = x.shape143    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)144    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]145    residual_scale_factor = b / sample_subset_size146    return brange, residual_scale_factor147 148 149def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):150    if scaling_vector is None:151        x_flat = x.flatten(1)152        residual = residual.flatten(1)153        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)154    else:155        x_plus_residual = scaled_index_add(156            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor157        )158    return x_plus_residual159 160 161attn_bias_cache: Dict[Tuple, Any] = {}162 163 164def get_attn_bias_and_cat(x_list, branges=None):165    """166    this will perform the index select, cat the tensors, and provide the attn_bias from cache167    """168    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]169    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))170    if all_shapes not in attn_bias_cache.keys():171        seqlens = []172        for b, x in zip(batch_sizes, x_list):173            for _ in range(b):174                seqlens.append(x.shape[1])175        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)176        attn_bias._batch_sizes = batch_sizes177        attn_bias_cache[all_shapes] = attn_bias178 179    if branges is not None:180        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])181    else:182        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)183        cat_tensors = torch.cat(tensors_bs1, dim=1)184 185    return attn_bias_cache[all_shapes], cat_tensors186 187 188def drop_add_residual_stochastic_depth_list(189    x_list: List[Tensor],190    residual_func: Callable[[Tensor, Any], Tensor],191    sample_drop_ratio: float = 0.0,192    scaling_vector=None,193) -> Tensor:194    # 1) generate random set of indices for dropping samples in the batch195    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]196    branges = [s[0] for s in branges_scales]197    residual_scale_factors = [s[1] for s in branges_scales]198 199    # 2) get attention bias and index+concat the tensors200    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)201 202    # 3) apply residual_func to get residual, and split the result203    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore204 205    outputs = []206    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):207        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))208    return outputs209 210 211class NestedTensorBlock(Block):212    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:213        """214        x_list contains a list of tensors to nest together and run215        """216        assert isinstance(self.attn, MemEffAttention)217 218        if self.training and self.sample_drop_ratio > 0.0:219 220            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:221                return self.attn(self.norm1(x), attn_bias=attn_bias)222 223            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:224                return self.mlp(self.norm2(x))225 226            x_list = drop_add_residual_stochastic_depth_list(227                x_list,228                residual_func=attn_residual_func,229                sample_drop_ratio=self.sample_drop_ratio,230                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,231            )232            x_list = drop_add_residual_stochastic_depth_list(233                x_list,234                residual_func=ffn_residual_func,235                sample_drop_ratio=self.sample_drop_ratio,236                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,237            )238            return x_list239        else:240 241            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:242                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))243 244            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:245                return self.ls2(self.mlp(self.norm2(x)))246 247            attn_bias, x = get_attn_bias_and_cat(x_list)248            x = x + attn_residual_func(x, attn_bias=attn_bias)249            x = x + ffn_residual_func(x)250            return attn_bias.split(x)251 252    def forward(self, x_or_x_list):253        if isinstance(x_or_x_list, Tensor):254            return super().forward(x_or_x_list)255        elif isinstance(x_or_x_list, list):256            if not XFORMERS_AVAILABLE:257                raise AssertionError("xFormers is required for using nested tensors")258            return self.forward_nested(x_or_x_list)259        else:260            raise AssertionError261