RuiTerrty/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/models/vision_transformer.py9 10import logging11import os12import warnings13 14from torch import Tensor15from torch import nn16 17 18logger = logging.getLogger("dinov2")19 20 21XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None22try:23    if XFORMERS_ENABLED:24        from xformers.ops import memory_efficient_attention, unbind25 26        XFORMERS_AVAILABLE = True27        warnings.warn("xFormers is available (Attention)")28    else:29        warnings.warn("xFormers is disabled (Attention)")30        raise ImportError31except ImportError:32    XFORMERS_AVAILABLE = False33    warnings.warn("xFormers is not available (Attention)")34 35 36class Attention(nn.Module):37    def __init__(38        self,39        dim: int,40        num_heads: int = 8,41        qkv_bias: bool = False,42        proj_bias: bool = True,43        attn_drop: float = 0.0,44        proj_drop: float = 0.0,45    ) -> None:46        super().__init__()47        self.num_heads = num_heads48        head_dim = dim // num_heads49        self.scale = head_dim**-0.550 51        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)52        self.attn_drop = nn.Dropout(attn_drop)53        self.proj = nn.Linear(dim, dim, bias=proj_bias)54        self.proj_drop = nn.Dropout(proj_drop)55 56    def forward(self, x: Tensor) -> Tensor:57        B, N, C = x.shape58        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)59 60        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]61        attn = q @ k.transpose(-2, -1)62 63        attn = attn.softmax(dim=-1)64        attn = self.attn_drop(attn)65 66        x = (attn @ v).transpose(1, 2).reshape(B, N, C)67        x = self.proj(x)68        x = self.proj_drop(x)69        return x70 71 72class MemEffAttention(Attention):73    def forward(self, x: Tensor, attn_bias=None) -> Tensor:74        if not XFORMERS_AVAILABLE:75            if attn_bias is not None:76                raise AssertionError("xFormers is required for using nested tensors")77            return super().forward(x)78 79        B, N, C = x.shape80        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)81 82        q, k, v = unbind(qkv, 2)83 84        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)85        x = x.reshape([B, N, C])86 87        x = self.proj(x)88        x = self.proj_drop(x)89        return x90