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 | |