Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from einops.einops import rearrange | |
| class Mlp(nn.Module): | |
| """Multi-Layer Perceptron (MLP)""" | |
| def __init__(self, | |
| in_dim, | |
| hidden_dim=None, | |
| out_dim=None, | |
| act_layer=nn.GELU): | |
| """ | |
| Args: | |
| in_dim: input features dimension | |
| hidden_dim: hidden features dimension | |
| out_dim: output features dimension | |
| act_layer: activation function | |
| """ | |
| super().__init__() | |
| out_dim = out_dim or in_dim | |
| hidden_dim = hidden_dim or in_dim | |
| self.fc1 = nn.Linear(in_dim, hidden_dim) | |
| self.act = act_layer() | |
| self.fc2 = nn.Linear(hidden_dim, out_dim) | |
| self.out_dim = out_dim | |
| def forward(self, x): | |
| x_size = x.size() | |
| x = x.view(-1, x_size[-1]) | |
| x = self.fc1(x) | |
| x = self.act(x) | |
| x = self.fc2(x) | |
| x = x.view(*x_size[:-1], self.out_dim) | |
| return x | |
| class VanillaAttention(nn.Module): | |
| def __init__(self, | |
| dim, | |
| num_heads=8, | |
| proj_bias=False): | |
| super().__init__() | |
| """ | |
| Args: | |
| dim: feature dimension | |
| num_heads: number of attention head | |
| proj_bias: bool use query, key, value bias | |
| """ | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.softmax_temp = self.head_dim ** -0.5 | |
| self.kv_proj = nn.Linear(dim, dim * 2, bias=proj_bias) | |
| self.q_proj = nn.Linear(dim, dim, bias=proj_bias) | |
| self.merge = nn.Linear(dim, dim) | |
| def forward(self, x_q, x_kv=None): | |
| """ | |
| Args: | |
| x_q (torch.Tensor): [N, L, C] | |
| x_kv (torch.Tensor): [N, S, C] | |
| """ | |
| if x_kv is None: | |
| x_kv = x_q | |
| bs, _, dim = x_q.shape | |
| bs, _, dim = x_kv.shape | |
| # [N, S, 2, H, D] => [2, N, H, S, D] | |
| kv = self.kv_proj(x_kv).reshape(bs, -1, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) | |
| # [N, L, H, D] => [N, H, L, D] | |
| q = self.q_proj(x_q).reshape(bs, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3) | |
| k, v = kv[0].transpose(-2, -1).contiguous(), kv[1].contiguous() # [N, H, D, S], [N, H, S, D] | |
| attn = (q @ k) * self.softmax_temp # [N, H, L, S] | |
| attn = attn.softmax(dim=-1) | |
| x_q = (attn @ v).transpose(1, 2).reshape(bs, -1, dim) | |
| x_q = self.merge(x_q) | |
| return x_q | |
| class CrossBidirectionalAttention(nn.Module): | |
| def __init__(self, dim, num_heads, proj_bias = False): | |
| super().__init__() | |
| """ | |
| Args: | |
| dim: feature dimension | |
| num_heads: number of attention head | |
| proj_bias: bool use query, key, value bias | |
| """ | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.softmax_temp = self.head_dim ** -0.5 | |
| self.qk_proj = nn.Linear(dim, dim, bias=proj_bias) | |
| self.v_proj = nn.Linear(dim, dim, bias=proj_bias) | |
| self.merge = nn.Linear(dim, dim, bias=proj_bias) | |
| self.temperature = nn.Parameter(torch.tensor([0.0]), requires_grad=True) | |
| # print(self.temperature) | |
| def map_(self, func, x0, x1): | |
| return func(x0), func(x1) | |
| def forward(self, x0, x1): | |
| """ | |
| Args: | |
| x0 (torch.Tensor): [N, L, C] | |
| x1 (torch.Tensor): [N, S, C] | |
| """ | |
| bs = x0.size(0) | |
| qk0, qk1 = self.map_(self.qk_proj, x0, x1) | |
| v0, v1 = self.map_(self.v_proj, x0, x1) | |
| qk0, qk1, v0, v1 = map( | |
| lambda t: t.reshape(bs, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous(), | |
| (qk0, qk1, v0, v1)) | |
| qk0, qk1 = qk0 * self.softmax_temp**0.5, qk1 * self.softmax_temp**0.5 | |
| sim = qk0 @ qk1.transpose(-2,-1).contiguous() | |
| attn01 = F.softmax(sim, dim=-1) | |
| attn10 = F.softmax(sim.transpose(-2, -1).contiguous(), dim=-1) | |
| x0 = attn01 @ v1 | |
| x1 = attn10 @ v0 | |
| x0, x1 = self.map_(lambda t: t.transpose(1, 2).flatten(start_dim=-2), | |
| x0, x1) | |
| x0, x1 = self.map_(self.merge, x0, x1) | |
| return x0, x1 | |
| class SwinPosEmbMLP(nn.Module): | |
| def __init__(self, | |
| dim): | |
| super().__init__() | |
| self.pos_embed = None | |
| self.pos_mlp = nn.Sequential(nn.Linear(2, 512, bias=True), | |
| nn.ReLU(), | |
| nn.Linear(512, dim, bias=False)) | |
| def forward(self, x): | |
| seq_length = x.shape[1] | |
| if self.pos_embed is None or self.training: | |
| seq_length = int(seq_length**0.5) | |
| coords = torch.arange(0, seq_length, device=x.device, dtype = x.dtype) | |
| grid = torch.stack(torch.meshgrid([coords, coords])).contiguous().unsqueeze(0) | |
| grid -= seq_length // 2 | |
| grid /= (seq_length // 2) | |
| self.pos_embed = self.pos_mlp(grid.flatten(2).transpose(1,2)) | |
| x = x + self.pos_embed | |
| return x | |
| class WindowSelfAttention(nn.Module): | |
| def __init__(self, dim, num_heads, mlp_hidden_coef, use_pre_pos_embed=False): | |
| super().__init__() | |
| self.mlp = Mlp(in_dim=dim*2, hidden_dim=dim*mlp_hidden_coef, out_dim=dim, act_layer=nn.GELU) | |
| self.gamma = nn.Parameter(torch.ones(dim)) | |
| self.norm1 = nn.LayerNorm(dim) | |
| self.norm2 = nn.LayerNorm(dim) | |
| self.attn = VanillaAttention(dim, num_heads=num_heads) | |
| self.pos_embed = SwinPosEmbMLP(dim) | |
| self.pos_embed_pre = SwinPosEmbMLP(dim) if use_pre_pos_embed else nn.Identity() | |
| def forward(self, x, x_pre): | |
| ww = x.shape[1] | |
| ww_pre = x_pre.shape[1] | |
| x = self.pos_embed(x) | |
| x_pre = self.pos_embed_pre(x_pre) | |
| x = torch.cat((x, x_pre), dim=1) | |
| x = x + self.gamma*self.norm1(self.mlp(torch.cat([x, self.attn(self.norm2(x))], dim=-1))) | |
| x, x_pre = x.split([ww, ww_pre], dim=1) | |
| return x, x_pre | |
| class WindowCrossAttention(nn.Module): | |
| def __init__(self, dim, num_heads, mlp_hidden_coef): | |
| super().__init__() | |
| self.norm1 = nn.LayerNorm(dim) | |
| self.norm2 = nn.LayerNorm(dim) | |
| self.mlp = Mlp(in_dim=dim*2, hidden_dim=dim*mlp_hidden_coef, out_dim=dim, act_layer=nn.GELU) | |
| self.cross_attn = CrossBidirectionalAttention(dim, num_heads=num_heads, proj_bias=False) | |
| self.gamma = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x0, x1): | |
| m_x0, m_x1 = self.cross_attn(self.norm1(x0), self.norm1(x1)) | |
| x0 = x0 + self.gamma*self.norm2(self.mlp(torch.cat([x0, m_x0], dim=-1))) | |
| x1 = x1 + self.gamma*self.norm2(self.mlp(torch.cat([x1, m_x1], dim=-1))) | |
| return x0, x1 | |
| class FineProcess(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| # Config | |
| block_dims = config['resnet']['block_dims'] | |
| self.block_dims = block_dims | |
| self.W_f = config['fine_window_size'] | |
| self.W_m = config['medium_window_size'] | |
| nhead_f = config["fine"]['nhead_fine_level'] | |
| nhead_m = config["fine"]['nhead_medium_level'] | |
| mlp_hidden_coef = config["fine"]['mlp_hidden_dim_coef'] | |
| # Networks | |
| self.conv_merge = nn.Sequential(nn.Conv2d(block_dims[2]*2, block_dims[1], kernel_size=1, stride=1, padding=0, bias=False), | |
| nn.Conv2d(block_dims[1], block_dims[1], kernel_size=3, stride=1, padding=1, groups=block_dims[1], bias=False), | |
| nn.BatchNorm2d(block_dims[1]) | |
| ) | |
| self.out_conv_m = nn.Conv2d(block_dims[1], block_dims[1], kernel_size=1, stride=1, padding=0, bias=False) | |
| self.out_conv_f = nn.Conv2d(block_dims[0], block_dims[0], kernel_size=1, stride=1, padding=0, bias=False) | |
| self.self_attn_m = WindowSelfAttention(block_dims[1], num_heads=nhead_m, | |
| mlp_hidden_coef=mlp_hidden_coef, use_pre_pos_embed=False) | |
| self.cross_attn_m = WindowCrossAttention(block_dims[1], num_heads=nhead_m, | |
| mlp_hidden_coef=mlp_hidden_coef) | |
| self.self_attn_f = WindowSelfAttention(block_dims[0], num_heads=nhead_f, | |
| mlp_hidden_coef=mlp_hidden_coef, use_pre_pos_embed=True) | |
| self.cross_attn_f = WindowCrossAttention(block_dims[0], num_heads=nhead_f, | |
| mlp_hidden_coef=mlp_hidden_coef) | |
| self.down_proj_m_f = nn.Linear(block_dims[1], block_dims[0], bias=False) | |
| for m in self.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') | |
| elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): | |
| nn.init.constant_(m.weight, 1) | |
| nn.init.constant_(m.bias, 0) | |
| def pre_process(self, feat_f0, feat_f1, feat_m0, feat_m1, feat_c0, feat_c1, feat_c0_pre, feat_c1_pre, data): | |
| W_f = self.W_f | |
| W_m = self.W_m | |
| data.update({'W_f': W_f, | |
| 'W_m': W_m}) | |
| # merge coarse features before and after loftr layer, and down proj channel dimesions | |
| feat_c0 = rearrange(feat_c0, 'n (h w) c -> n c h w', h =data["hw0_c"][0], w =data["hw0_c"][1]) | |
| feat_c1 = rearrange(feat_c1, 'n (h w) c -> n c h w', h =data["hw1_c"][0], w =data["hw1_c"][1]) | |
| feat_c0 = self.conv_merge(torch.cat([feat_c0, feat_c0_pre], dim=1)) | |
| feat_c1 = self.conv_merge(torch.cat([feat_c1, feat_c1_pre], dim=1)) | |
| feat_c0 = rearrange(feat_c0, 'n c h w -> n (h w) 1 c') | |
| feat_c1 = rearrange(feat_c1, 'n c h w -> n (h w) 1 c') | |
| stride_f = data['hw0_f'][0] // data['hw0_c'][0] | |
| stride_m = data['hw0_m'][0] // data['hw0_c'][0] | |
| if feat_m0.shape[2] == feat_m1.shape[2] and feat_m0.shape[3] == feat_m1.shape[3]: | |
| feat_m = self.out_conv_m(torch.cat([feat_m0, feat_m1], dim=0)) | |
| feat_m0, feat_m1 = torch.chunk(feat_m, 2, dim=0) | |
| feat_f = self.out_conv_f(torch.cat([feat_f0, feat_f1], dim=0)) | |
| feat_f0, feat_f1 = torch.chunk(feat_f, 2, dim=0) | |
| else: | |
| feat_m0 = self.out_conv_m(feat_m0) | |
| feat_m1 = self.out_conv_m(feat_m1) | |
| feat_f0 = self.out_conv_f(feat_f0) | |
| feat_f1 = self.out_conv_f(feat_f1) | |
| # 1. unfold (crop windows) all local windows | |
| feat_m0_unfold = F.unfold(feat_m0, kernel_size=(W_m, W_m), stride=stride_m, padding=W_m//2) | |
| feat_m0_unfold = rearrange(feat_m0_unfold, 'n (c ww) l -> n l ww c', ww=W_m**2) | |
| feat_m1_unfold = F.unfold(feat_m1, kernel_size=(W_m, W_m), stride=stride_m, padding=W_m//2) | |
| feat_m1_unfold = rearrange(feat_m1_unfold, 'n (c ww) l -> n l ww c', ww=W_m**2) | |
| feat_f0_unfold = F.unfold(feat_f0, kernel_size=(W_f, W_f), stride=stride_f, padding=W_f//2) | |
| feat_f0_unfold = rearrange(feat_f0_unfold, 'n (c ww) l -> n l ww c', ww=W_f**2) | |
| feat_f1_unfold = F.unfold(feat_f1, kernel_size=(W_f, W_f), stride=stride_f, padding=W_f//2) | |
| feat_f1_unfold = rearrange(feat_f1_unfold, 'n (c ww) l -> n l ww c', ww=W_f**2) | |
| # 2. select only the predicted matches | |
| feat_c0 = feat_c0[data['b_ids'], data['i_ids']] # [n, ww, cm] | |
| feat_c1 = feat_c1[data['b_ids'], data['j_ids']] | |
| feat_m0_unfold = feat_m0_unfold[data['b_ids'], data['i_ids']] # [n, ww, cm] | |
| feat_m1_unfold = feat_m1_unfold[data['b_ids'], data['j_ids']] | |
| feat_f0_unfold = feat_f0_unfold[data['b_ids'], data['i_ids']] # [n, ww, cf] | |
| feat_f1_unfold = feat_f1_unfold[data['b_ids'], data['j_ids']] | |
| return feat_c0, feat_c1, feat_m0_unfold, feat_m1_unfold, feat_f0_unfold, feat_f1_unfold | |
| def forward(self, feat_f0, feat_f1, feat_m0, feat_m1, feat_c0, feat_c1, feat_c0_pre, feat_c1_pre, data): | |
| """ | |
| Args: | |
| feat_f0 (torch.Tensor): [N, C, H, W] | |
| feat_f1 (torch.Tensor): [N, C, H, W] | |
| feat_m0 (torch.Tensor): [N, C, H, W] | |
| feat_m1 (torch.Tensor): [N, C, H, W] | |
| feat_c0 (torch.Tensor): [N, L, C] | |
| feat_c1 (torch.Tensor): [N, S, C] | |
| feat_c0_pre (torch.Tensor): [N, C, H, W] | |
| feat_c1_pre (torch.Tensor): [N, C, H, W] | |
| data (dict): with keys ['hw0_c', 'hw1_c', 'hw0_m', 'hw1_m', 'hw0_f', 'hw1_f', 'b_ids', 'j_ids'] | |
| """ | |
| # TODO: Check for this case | |
| if data['b_ids'].shape[0] == 0: | |
| feat0 = torch.empty(0, self.W_f**2, self.block_dims[0], device=feat_f0.device) | |
| feat1 = torch.empty(0, self.W_f**2, self.block_dims[0], device=feat_f0.device) | |
| return feat0, feat1 | |
| feat_c0, feat_c1, feat_m0_unfold, feat_m1_unfold, \ | |
| feat_f0_unfold, feat_f1_unfold = self.pre_process(feat_f0, feat_f1, feat_m0, feat_m1, | |
| feat_c0, feat_c1, feat_c0_pre, feat_c1_pre, data) | |
| # self attention (c + m) | |
| feat_m_unfold, _ = self.self_attn_m(torch.cat([feat_m0_unfold, feat_m1_unfold], dim=0), | |
| torch.cat([feat_c0, feat_c1], dim=0)) | |
| feat_m0_unfold, feat_m1_unfold = torch.chunk(feat_m_unfold, 2, dim=0) | |
| # cross attention (m0 <-> m1) | |
| feat_m0_unfold, feat_m1_unfold = self.cross_attn_m(feat_m0_unfold, feat_m1_unfold) | |
| # down proj m | |
| feat_m_unfold = self.down_proj_m_f(torch.cat([feat_m0_unfold, feat_m1_unfold], dim=0)) | |
| feat_m0_unfold, feat_m1_unfold = torch.chunk(feat_m_unfold, 2, dim=0) | |
| # self attention (m + f) | |
| feat_f_unfold, _ = self.self_attn_f(torch.cat([feat_f0_unfold, feat_f1_unfold], dim=0), | |
| torch.cat([feat_m0_unfold, feat_m1_unfold], dim=0)) | |
| feat_f0_unfold, feat_f1_unfold = torch.chunk(feat_f_unfold, 2, dim=0) | |
| # cross attention (f0 <-> f1) | |
| feat_f0_unfold, feat_f1_unfold = self.cross_attn_f(feat_f0_unfold, feat_f1_unfold) | |
| return feat_f0_unfold, feat_f1_unfold | |