naver/SuperFeatures

1# Copyright (C) 2021-2022 Naver Corporation. All rights reserved.
2# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
3
4import torch
5from torch import nn
6
7class LocalfeatureIntegrationTransformer(nn.Module):
8    """Map a set of local features to a fixed number of SuperFeatures """
9
10    def __init__(self, T, N, input_dim, dim):
11        """
12        T: number of iterations
13        N: number of SuperFeatures
14        input_dim: dimension of input local features
15        dim: dimension of SuperFeatures
16        """
17        super().__init__()
18        self.T = T
19        self.N = N
20        self.input_dim = input_dim
21        self.dim = dim
22        # learnable initialization
23        self.templates_init = nn.Parameter(torch.randn(1,self.N,dim))
24        # qkv
25        self.project_q = nn.Linear(dim, dim, bias=False)
26        self.project_k = nn.Linear(input_dim, dim, bias=False)
27        self.project_v = nn.Linear(input_dim, dim, bias=False)
28        # layer norms
29        self.norm_inputs = nn.LayerNorm(input_dim)
30        self.norm_templates = nn.LayerNorm(dim)
31        # for the normalization
32        self.softmax = nn.Softmax(dim=-1)
33        self.scale = dim ** -0.5
34        # mlp
35        self.norm_mlp = nn.LayerNorm(dim)
36        mlp_dim = dim//2
37        self.mlp = nn.Sequential(nn.Linear(dim, mlp_dim), nn.ReLU(), nn.Linear(mlp_dim, dim) )
38
39
40    def forward(self, x):
41        """
42        input:
43            x has shape BxCxHxW
44        output:
45            template (output SuperFeatures): tensor of shape BxCxNx1
46            attn (attention over local features at the last iteration): tensor of shape BxNxHxW
47        """
48        # reshape inputs from BxCxHxW to Bx(H*W)xC
49        B,C,H,W = x.size()
50        x = x.reshape(B,C,H*W).permute(0,2,1)
51
52        # k and v projection
53        x = self.norm_inputs(x)
54        k = self.project_k(x)
55        v = self.project_v(x)
56
57        # template initialization
58        templates = torch.repeat_interleave(self.templates_init, B, dim=0)
59        attn = None
60
61        # main iteration loop
62        for _ in range(self.T):
63            templates_prev = templates
64
65            # q projection
66            templates = self.norm_templates(templates)
67            q = self.project_q(templates)
68
69            # attention
70            q = q * self.scale  # Normalization.
71            attn_logits =  torch.einsum('bnd,bld->bln', q, k)
72            attn = self.softmax(attn_logits)
73            attn = attn + 1e-8 # to avoid zero when with the L1 norm below
74            attn = attn / attn.sum(dim=-2, keepdim=True)
75
76            # update template
77            templates = templates_prev + torch.einsum('bld,bln->bnd', v, attn)
78
79            # mlp
80            templates = templates + self.mlp(self.norm_mlp(templates))
81
82        # reshape templates to BxDxNx1
83        templates = templates.permute(0,2,1)[:,:,:,None]
84        attn = attn.permute(0,2,1).view(B,self.N,H,W)
85
86        return templates, attn
87
88    def __repr__(self):
89        s = str(self.__class__.__name__)
90        for k in ["T","N","input_dim","dim"]:
91            s += "\n  {:s}: {:d}".format(k, getattr(self,k))
92        return s
93