cross_bottleneck_attn

models/cross_bottleneck_attn.py

@@ -0,0 +1,116 @@
+"""
+Code based on timm https://github.com/huggingface/pytorch-image-models
+
+Modifications and additions for mivolo by / Copyright 2023, Irina Tolstykh, Maxim Kuprashevich
+"""
+
+import torch
+import torch.nn as nn
+from timm.layers.bottleneck_attn import PosEmbedRel
+from timm.layers.helpers import make_divisible
+from timm.layers.mlp import Mlp
+from timm.layers.trace_utils import _assert
+from timm.layers.weight_init import trunc_normal_
+
+
+class CrossBottleneckAttn(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out=None,
+        feat_size=None,
+        stride=1,
+        num_heads=4,
+        dim_head=None,
+        qk_ratio=1.0,
+        qkv_bias=False,
+        scale_pos_embed=False,
+    ):
+        super().__init__()
+        assert feat_size is not None, "A concrete feature size matching expected input (H, W) is required"
+        dim_out = dim_out or dim
+        assert dim_out % num_heads == 0
+
+        self.num_heads = num_heads
+        self.dim_head_qk = dim_head or make_divisible(dim_out * qk_ratio, divisor=8) // num_heads
+        self.dim_head_v = dim_out // self.num_heads
+        self.dim_out_qk = num_heads * self.dim_head_qk
+        self.dim_out_v = num_heads * self.dim_head_v
+        self.scale = self.dim_head_qk**-0.5
+        self.scale_pos_embed = scale_pos_embed
+
+        self.qkv_f = nn.Conv2d(dim, self.dim_out_qk * 2 + self.dim_out_v, 1, bias=qkv_bias)
+        self.qkv_p = nn.Conv2d(dim, self.dim_out_qk * 2 + self.dim_out_v, 1, bias=qkv_bias)
+
+        # NOTE I'm only supporting relative pos embedding for now
+        self.pos_embed = PosEmbedRel(feat_size, dim_head=self.dim_head_qk, scale=self.scale)
+
+        self.norm = nn.LayerNorm([self.dim_out_v * 2, *feat_size])
+        mlp_ratio = 4
+        self.mlp = Mlp(
+            in_features=self.dim_out_v * 2,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=nn.GELU,
+            out_features=dim_out,
+            drop=0,
+            use_conv=True,
+        )
+
+        self.pool = nn.AvgPool2d(2, 2) if stride == 2 else nn.Identity()
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        trunc_normal_(self.qkv_f.weight, std=self.qkv_f.weight.shape[1] ** -0.5)  # fan-in
+        trunc_normal_(self.qkv_p.weight, std=self.qkv_p.weight.shape[1] ** -0.5)  # fan-in
+        trunc_normal_(self.pos_embed.height_rel, std=self.scale)
+        trunc_normal_(self.pos_embed.width_rel, std=self.scale)
+
+    def get_qkv(self, x, qvk_conv):
+        B, C, H, W = x.shape
+
+        x = qvk_conv(x)  # B, (2 * dim_head_qk + dim_head_v) * num_heads, H, W
+
+        q, k, v = torch.split(x, [self.dim_out_qk, self.dim_out_qk, self.dim_out_v], dim=1)
+
+        q = q.reshape(B * self.num_heads, self.dim_head_qk, -1).transpose(-1, -2)
+        k = k.reshape(B * self.num_heads, self.dim_head_qk, -1)  # no transpose, for q @ k
+        v = v.reshape(B * self.num_heads, self.dim_head_v, -1).transpose(-1, -2)
+
+        return q, k, v
+
+    def apply_attn(self, q, k, v, B, H, W, dropout=None):
+        if self.scale_pos_embed:
+            attn = (q @ k + self.pos_embed(q)) * self.scale  # B * num_heads, H * W, H * W
+        else:
+            attn = (q @ k) * self.scale + self.pos_embed(q)
+        attn = attn.softmax(dim=-1)
+        if dropout:
+            attn = dropout(attn)
+
+        out = (attn @ v).transpose(-1, -2).reshape(B, self.dim_out_v, H, W)  # B, dim_out, H, W
+        return out
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+
+        dim = int(C / 2)
+        x1 = x[:, :dim, :, :]
+        x2 = x[:, dim:, :, :]
+
+        _assert(H == self.pos_embed.height, "")
+        _assert(W == self.pos_embed.width, "")
+
+        q_f, k_f, v_f = self.get_qkv(x1, self.qkv_f)
+        q_p, k_p, v_p = self.get_qkv(x2, self.qkv_p)
+
+        # person to face
+        out_f = self.apply_attn(q_f, k_p, v_p, B, H, W)
+        # face to person
+        out_p = self.apply_attn(q_p, k_f, v_f, B, H, W)
+
+        x_pf = torch.cat((out_f, out_p), dim=1)  # B, dim_out * 2, H, W
+        x_pf = self.norm(x_pf)
+        x_pf = self.mlp(x_pf)  # B, dim_out, H, W
+
+        out = self.pool(x_pf)
+        return out