Vincentqyw
fix: roma
c74a070
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}