Spaces:
Build error
Build error
| # Copyright (C) 2021-2022 Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| import torch | |
| from torch import nn | |
| class LocalfeatureIntegrationTransformer(nn.Module): | |
| """Map a set of local features to a fixed number of SuperFeatures """ | |
| def __init__(self, T, N, input_dim, dim): | |
| """ | |
| T: number of iterations | |
| N: number of SuperFeatures | |
| input_dim: dimension of input local features | |
| dim: dimension of SuperFeatures | |
| """ | |
| super().__init__() | |
| self.T = T | |
| self.N = N | |
| self.input_dim = input_dim | |
| self.dim = dim | |
| # learnable initialization | |
| self.templates_init = nn.Parameter(torch.randn(1,self.N,dim)) | |
| # qkv | |
| self.project_q = nn.Linear(dim, dim, bias=False) | |
| self.project_k = nn.Linear(input_dim, dim, bias=False) | |
| self.project_v = nn.Linear(input_dim, dim, bias=False) | |
| # layer norms | |
| self.norm_inputs = nn.LayerNorm(input_dim) | |
| self.norm_templates = nn.LayerNorm(dim) | |
| # for the normalization | |
| self.softmax = nn.Softmax(dim=-1) | |
| self.scale = dim ** -0.5 | |
| # mlp | |
| self.norm_mlp = nn.LayerNorm(dim) | |
| mlp_dim = dim//2 | |
| self.mlp = nn.Sequential(nn.Linear(dim, mlp_dim), nn.ReLU(), nn.Linear(mlp_dim, dim) ) | |
| def forward(self, x): | |
| """ | |
| input: | |
| x has shape BxCxHxW | |
| output: | |
| template (output SuperFeatures): tensor of shape BxCxNx1 | |
| attn (attention over local features at the last iteration): tensor of shape BxNxHxW | |
| """ | |
| # reshape inputs from BxCxHxW to Bx(H*W)xC | |
| B,C,H,W = x.size() | |
| x = x.reshape(B,C,H*W).permute(0,2,1) | |
| # k and v projection | |
| x = self.norm_inputs(x) | |
| k = self.project_k(x) | |
| v = self.project_v(x) | |
| # template initialization | |
| templates = torch.repeat_interleave(self.templates_init, B, dim=0) | |
| attn = None | |
| # main iteration loop | |
| for _ in range(self.T): | |
| templates_prev = templates | |
| # q projection | |
| templates = self.norm_templates(templates) | |
| q = self.project_q(templates) | |
| # attention | |
| q = q * self.scale # Normalization. | |
| attn_logits = torch.einsum('bnd,bld->bln', q, k) | |
| attn = self.softmax(attn_logits) | |
| attn = attn + 1e-8 # to avoid zero when with the L1 norm below | |
| attn = attn / attn.sum(dim=-2, keepdim=True) | |
| # update template | |
| templates = templates_prev + torch.einsum('bld,bln->bnd', v, attn) | |
| # mlp | |
| templates = templates + self.mlp(self.norm_mlp(templates)) | |
| # reshape templates to BxDxNx1 | |
| templates = templates.permute(0,2,1)[:,:,:,None] | |
| attn = attn.permute(0,2,1).view(B,self.N,H,W) | |
| return templates, attn | |
| def __repr__(self): | |
| s = str(self.__class__.__name__) | |
| for k in ["T","N","input_dim","dim"]: | |
| s += "\n {:s}: {:d}".format(k, getattr(self,k)) | |
| return s | |