import torch import torch.nn as nn import time eps = 1e-8 def sinkhorn(M, r, c, iteration): p = torch.softmax(M, dim=-1) u = torch.ones_like(r) v = torch.ones_like(c) for _ in range(iteration): u = r / ((p * v.unsqueeze(-2)).sum(-1) + eps) v = c / ((p * u.unsqueeze(-1)).sum(-2) + eps) p = p * u.unsqueeze(-1) * v.unsqueeze(-2) return p def sink_algorithm(M, dustbin, iteration): M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1) M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2) r = torch.ones([M.shape[0], M.shape[1] - 1], device="cuda") r = torch.cat([r, torch.ones([M.shape[0], 1], device="cuda") * M.shape[1]], dim=-1) c = torch.ones([M.shape[0], M.shape[2] - 1], device="cuda") c = torch.cat([c, torch.ones([M.shape[0], 1], device="cuda") * M.shape[2]], dim=-1) p = sinkhorn(M, r, c, iteration) return p class attention_block(nn.Module): def __init__(self, channels, head, type): assert type == "self" or type == "cross", "invalid attention type" nn.Module.__init__(self) self.head = head self.type = type self.head_dim = channels // head self.query_filter = nn.Conv1d(channels, channels, kernel_size=1) self.key_filter = nn.Conv1d(channels, channels, kernel_size=1) self.value_filter = nn.Conv1d(channels, channels, kernel_size=1) self.attention_filter = nn.Sequential( nn.Conv1d(2 * channels, 2 * channels, kernel_size=1), nn.SyncBatchNorm(2 * channels), nn.ReLU(), nn.Conv1d(2 * channels, channels, kernel_size=1), ) self.mh_filter = nn.Conv1d(channels, channels, kernel_size=1) def forward(self, fea1, fea2): batch_size, n, m = fea1.shape[0], fea1.shape[2], fea2.shape[2] query1, key1, value1 = ( self.query_filter(fea1).view(batch_size, self.head_dim, self.head, -1), self.key_filter(fea1).view(batch_size, self.head_dim, self.head, -1), self.value_filter(fea1).view(batch_size, self.head_dim, self.head, -1), ) query2, key2, value2 = ( self.query_filter(fea2).view(batch_size, self.head_dim, self.head, -1), self.key_filter(fea2).view(batch_size, self.head_dim, self.head, -1), self.value_filter(fea2).view(batch_size, self.head_dim, self.head, -1), ) if self.type == "self": score1, score2 = torch.softmax( torch.einsum("bdhn,bdhm->bhnm", query1, key1) / self.head_dim**0.5, dim=-1, ), torch.softmax( torch.einsum("bdhn,bdhm->bhnm", query2, key2) / self.head_dim**0.5, dim=-1, ) add_value1, add_value2 = torch.einsum( "bhnm,bdhm->bdhn", score1, value1 ), torch.einsum("bhnm,bdhm->bdhn", score2, value2) else: score1, score2 = torch.softmax( torch.einsum("bdhn,bdhm->bhnm", query1, key2) / self.head_dim**0.5, dim=-1, ), torch.softmax( torch.einsum("bdhn,bdhm->bhnm", query2, key1) / self.head_dim**0.5, dim=-1, ) add_value1, add_value2 = torch.einsum( "bhnm,bdhm->bdhn", score1, value2 ), torch.einsum("bhnm,bdhm->bdhn", score2, value1) add_value1, add_value2 = self.mh_filter( add_value1.contiguous().view(batch_size, self.head * self.head_dim, n) ), self.mh_filter( add_value2.contiguous().view(batch_size, self.head * self.head_dim, m) ) fea11, fea22 = torch.cat([fea1, add_value1], dim=1), torch.cat( [fea2, add_value2], dim=1 ) fea1, fea2 = fea1 + self.attention_filter(fea11), fea2 + self.attention_filter( fea22 ) return fea1, fea2 class matcher(nn.Module): def __init__(self, config): nn.Module.__init__(self) self.use_score_encoding = config.use_score_encoding self.layer_num = config.layer_num self.sink_iter = config.sink_iter self.position_encoder = nn.Sequential( nn.Conv1d(3, 32, kernel_size=1) if config.use_score_encoding else nn.Conv1d(2, 32, kernel_size=1), nn.SyncBatchNorm(32), nn.ReLU(), nn.Conv1d(32, 64, kernel_size=1), nn.SyncBatchNorm(64), nn.ReLU(), nn.Conv1d(64, 128, kernel_size=1), nn.SyncBatchNorm(128), nn.ReLU(), nn.Conv1d(128, 256, kernel_size=1), nn.SyncBatchNorm(256), nn.ReLU(), nn.Conv1d(256, config.net_channels, kernel_size=1), ) self.dustbin = nn.Parameter(torch.tensor(1, dtype=torch.float32, device="cuda")) self.self_attention_block = nn.Sequential( *[ attention_block(config.net_channels, config.head, "self") for _ in range(config.layer_num) ] ) self.cross_attention_block = nn.Sequential( *[ attention_block(config.net_channels, config.head, "cross") for _ in range(config.layer_num) ] ) self.final_project = nn.Conv1d( config.net_channels, config.net_channels, kernel_size=1 ) def forward(self, data, test_mode=True): desc1, desc2 = data["desc1"], data["desc2"] desc1, desc2 = torch.nn.functional.normalize( desc1, dim=-1 ), torch.nn.functional.normalize(desc2, dim=-1) desc1, desc2 = desc1.transpose(1, 2), desc2.transpose(1, 2) if test_mode: encode_x1, encode_x2 = data["x1"], data["x2"] else: encode_x1, encode_x2 = data["aug_x1"], data["aug_x2"] if not self.use_score_encoding: encode_x1, encode_x2 = encode_x1[:, :, :2], encode_x2[:, :, :2] encode_x1, encode_x2 = encode_x1.transpose(1, 2), encode_x2.transpose(1, 2) x1_pos_embedding, x2_pos_embedding = self.position_encoder( encode_x1 ), self.position_encoder(encode_x2) aug_desc1, aug_desc2 = x1_pos_embedding + desc1, x2_pos_embedding + desc2 for i in range(self.layer_num): aug_desc1, aug_desc2 = self.self_attention_block[i](aug_desc1, aug_desc2) aug_desc1, aug_desc2 = self.cross_attention_block[i](aug_desc1, aug_desc2) aug_desc1, aug_desc2 = self.final_project(aug_desc1), self.final_project( aug_desc2 ) desc_mat = torch.matmul(aug_desc1.transpose(1, 2), aug_desc2) p = sink_algorithm(desc_mat, self.dustbin, self.sink_iter[0]) return {"p": p}