Spaces:
Build error
Build error
File size: 3,210 Bytes
9651aac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# 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
|