| import logging
|
| import os
|
| import warnings
|
|
|
| import torch
|
| from torch import Tensor, nn
|
| from torch.nn import init
|
| from torch.nn.functional import softmax
|
|
|
| logger = logging.getLogger("dinov2")
|
|
|
|
|
| XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
|
| try:
|
| if XFORMERS_ENABLED:
|
| from xformers.ops import memory_efficient_attention, unbind
|
|
|
| XFORMERS_AVAILABLE = True
|
| else:
|
| warnings.warn("xFormers is disabled (Attention)")
|
| raise ImportError
|
| except ImportError:
|
| XFORMERS_AVAILABLE = False
|
| warnings.warn("xFormers is not available (Attention)")
|
|
|
|
|
| class Attention(nn.Module):
|
| def __init__(
|
| self,
|
| dim: int,
|
| num_heads: int = 8,
|
| qkv_bias: bool = False,
|
| proj_bias: bool = True,
|
| attn_drop: float = 0.0,
|
| proj_drop: float = 0.0,
|
| ) -> None:
|
| super().__init__()
|
| self.num_heads = num_heads
|
| head_dim = dim // num_heads
|
| self.scale = head_dim**-0.5
|
|
|
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
| self.attn_drop = nn.Dropout(attn_drop)
|
| self.proj = nn.Linear(dim, dim, bias=proj_bias)
|
| self.proj_drop = nn.Dropout(proj_drop)
|
|
|
| def forward(self, x: Tensor) -> Tensor:
|
| B, N, C = x.shape
|
| qkv = (
|
| self.qkv(x)
|
| .reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
| .permute(2, 0, 3, 1, 4)
|
| )
|
|
|
| q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
|
| attn = q @ k.transpose(-2, -1)
|
|
|
| attn = attn.softmax(dim=-1)
|
| attn = self.attn_drop(attn)
|
|
|
| x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
| x = self.proj(x)
|
| x = self.proj_drop(x)
|
| return x
|
|
|
|
|
| class MemEffAttention(Attention):
|
| def forward(self, x: Tensor, attn_bias=None) -> Tensor:
|
| if not XFORMERS_AVAILABLE:
|
| if attn_bias is not None:
|
| raise AssertionError("xFormers is required for using nested tensors")
|
| return super().forward(x)
|
|
|
| B, N, C = x.shape
|
| qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
|
|
| q, k, v = unbind(qkv, 2)
|
|
|
| x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
|
| x = x.reshape([B, N, C])
|
|
|
| x = self.proj(x)
|
| x = self.proj_drop(x)
|
| return x
|
|
|
|
|
| class AttentionBlock(nn.Module):
|
| def __init__(self, state_dim):
|
| super().__init__()
|
|
|
| q = nn.Parameter(torch.randn(1, state_dim, 1))
|
|
|
| init.xavier_uniform_(q)
|
|
|
| self.q = q
|
|
|
| self.state_dim = state_dim
|
|
|
| def forward(self, x):
|
| x = x.view(tuple(x.shape[:2]) + (-1,))
|
|
|
| r = self.q * x
|
|
|
| unnormalized_attn_map = r.sum(1)
|
|
|
| attn_map = softmax(unnormalized_attn_map, dim=1)
|
|
|
| a = (torch.unsqueeze(attn_map, 1) * x).sum(2)
|
|
|
| return a
|
|
|
