Spaces:
Running
Running
| from loguru import logger | |
| import copy | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .linear_attention import LinearAttention, FullAttention | |
| class LoFTREncoderLayer(nn.Module): | |
| def __init__(self, d_model, nhead, attention="linear"): | |
| super(LoFTREncoderLayer, self).__init__() | |
| self.dim = d_model // nhead | |
| self.nhead = nhead | |
| # multi-head attention | |
| self.q_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.k_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.v_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.attention = LinearAttention() if attention == "linear" else FullAttention() | |
| self.merge = nn.Linear(d_model, d_model, bias=False) | |
| # feed-forward network | |
| self.mlp = nn.Sequential( | |
| nn.Linear(d_model * 2, d_model * 2, bias=False), | |
| nn.GELU(), | |
| nn.Linear(d_model * 2, d_model, bias=False), | |
| ) | |
| # norm and dropout | |
| self.norm1 = nn.LayerNorm(d_model) | |
| self.norm2 = nn.LayerNorm(d_model) | |
| def forward(self, x, source, x_mask=None, source_mask=None): | |
| """ | |
| Args: | |
| x (torch.Tensor): [N, L, C] | |
| source (torch.Tensor): [N, S, C] | |
| x_mask (torch.Tensor): [N, L] (optional) | |
| source_mask (torch.Tensor): [N, S] (optional) | |
| """ | |
| bs = x.shape[0] | |
| query, key, value = x, source, source | |
| # multi-head attention | |
| query = self.q_proj(query).view(bs, -1, self.nhead, self.dim) # [N, L, (H, D)] | |
| key = self.k_proj(key).view(bs, -1, self.nhead, self.dim) # [N, S, (H, D)] | |
| value = self.v_proj(value).view(bs, -1, self.nhead, self.dim) | |
| message = self.attention( | |
| query, key, value, q_mask=x_mask, kv_mask=source_mask | |
| ) # [N, L, (H, D)] | |
| message = self.merge(message.view(bs, -1, self.nhead * self.dim)) # [N, L, C] | |
| message = self.norm1(message) | |
| # feed-forward network | |
| message = self.mlp(torch.cat([x, message], dim=2)) | |
| message = self.norm2(message) | |
| return x + message | |
| class TopicFormer(nn.Module): | |
| """A Local Feature Transformer (LoFTR) module.""" | |
| def __init__(self, config): | |
| super(TopicFormer, self).__init__() | |
| self.config = config | |
| self.d_model = config["d_model"] | |
| self.nhead = config["nhead"] | |
| self.layer_names = config["layer_names"] | |
| encoder_layer = LoFTREncoderLayer( | |
| config["d_model"], config["nhead"], config["attention"] | |
| ) | |
| self.layers = nn.ModuleList( | |
| [copy.deepcopy(encoder_layer) for _ in range(len(self.layer_names))] | |
| ) | |
| self.topic_transformers = ( | |
| nn.ModuleList( | |
| [ | |
| copy.deepcopy(encoder_layer) | |
| for _ in range(2 * config["n_topic_transformers"]) | |
| ] | |
| ) | |
| if config["n_samples"] > 0 | |
| else None | |
| ) # nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(2)]) | |
| self.n_iter_topic_transformer = config["n_topic_transformers"] | |
| self.seed_tokens = nn.Parameter( | |
| torch.randn(config["n_topics"], config["d_model"]) | |
| ) | |
| self.register_parameter("seed_tokens", self.seed_tokens) | |
| self.n_samples = config["n_samples"] | |
| self._reset_parameters() | |
| def _reset_parameters(self): | |
| for p in self.parameters(): | |
| if p.dim() > 1: | |
| nn.init.xavier_uniform_(p) | |
| def sample_topic(self, prob_topics, topics, L): | |
| """ | |
| Args: | |
| topics (torch.Tensor): [N, L+S, K] | |
| """ | |
| prob_topics0, prob_topics1 = prob_topics[:, :L], prob_topics[:, L:] | |
| topics0, topics1 = topics[:, :L], topics[:, L:] | |
| theta0 = F.normalize(prob_topics0.sum(dim=1), p=1, dim=-1) # [N, K] | |
| theta1 = F.normalize(prob_topics1.sum(dim=1), p=1, dim=-1) | |
| theta = F.normalize(theta0 * theta1, p=1, dim=-1) | |
| if self.n_samples == 0: | |
| return None | |
| if self.training: | |
| sampled_inds = torch.multinomial(theta, self.n_samples) | |
| sampled_values = torch.gather(theta, dim=-1, index=sampled_inds) | |
| else: | |
| sampled_values, sampled_inds = torch.topk(theta, self.n_samples, dim=-1) | |
| sampled_topics0 = torch.gather( | |
| topics0, | |
| dim=-1, | |
| index=sampled_inds.unsqueeze(1).repeat(1, topics0.shape[1], 1), | |
| ) | |
| sampled_topics1 = torch.gather( | |
| topics1, | |
| dim=-1, | |
| index=sampled_inds.unsqueeze(1).repeat(1, topics1.shape[1], 1), | |
| ) | |
| return sampled_topics0, sampled_topics1 | |
| def reduce_feat(self, feat, topick, N, C): | |
| len_topic = topick.sum(dim=-1).int() | |
| max_len = len_topic.max().item() | |
| selected_ids = topick.bool() | |
| resized_feat = torch.zeros( | |
| (N, max_len, C), dtype=torch.float32, device=feat.device | |
| ) | |
| new_mask = torch.zeros_like(resized_feat[..., 0]).bool() | |
| for i in range(N): | |
| new_mask[i, : len_topic[i]] = True | |
| resized_feat[new_mask, :] = feat[selected_ids, :] | |
| return resized_feat, new_mask, selected_ids | |
| def forward(self, feat0, feat1, mask0=None, mask1=None): | |
| """ | |
| Args: | |
| feat0 (torch.Tensor): [N, L, C] | |
| feat1 (torch.Tensor): [N, S, C] | |
| mask0 (torch.Tensor): [N, L] (optional) | |
| mask1 (torch.Tensor): [N, S] (optional) | |
| """ | |
| assert ( | |
| self.d_model == feat0.shape[2] | |
| ), "the feature number of src and transformer must be equal" | |
| N, L, S, C, K = ( | |
| feat0.shape[0], | |
| feat0.shape[1], | |
| feat1.shape[1], | |
| feat0.shape[2], | |
| self.config["n_topics"], | |
| ) | |
| seeds = self.seed_tokens.unsqueeze(0).repeat(N, 1, 1) | |
| feat = torch.cat((feat0, feat1), dim=1) | |
| if mask0 is not None: | |
| mask = torch.cat((mask0, mask1), dim=-1) | |
| else: | |
| mask = None | |
| for layer, name in zip(self.layers, self.layer_names): | |
| if name == "seed": | |
| # seeds = layer(seeds, feat0, None, mask0) | |
| # seeds = layer(seeds, feat1, None, mask1) | |
| seeds = layer(seeds, feat, None, mask) | |
| elif name == "feat": | |
| feat0 = layer(feat0, seeds, mask0, None) | |
| feat1 = layer(feat1, seeds, mask1, None) | |
| dmatrix = torch.einsum("nmd,nkd->nmk", feat, seeds) | |
| prob_topics = F.softmax(dmatrix, dim=-1) | |
| feat_topics = torch.zeros_like(dmatrix).scatter_( | |
| -1, torch.argmax(dmatrix, dim=-1, keepdim=True), 1.0 | |
| ) | |
| if mask is not None: | |
| feat_topics = feat_topics * mask.unsqueeze(-1) | |
| prob_topics = prob_topics * mask.unsqueeze(-1) | |
| if (feat_topics.detach().sum(dim=1).sum(dim=0) > 100).sum() <= 3: | |
| logger.warning("topic distribution is highly sparse!") | |
| sampled_topics = self.sample_topic(prob_topics.detach(), feat_topics, L) | |
| if sampled_topics is not None: | |
| updated_feat0, updated_feat1 = torch.zeros_like(feat0), torch.zeros_like( | |
| feat1 | |
| ) | |
| s_topics0, s_topics1 = sampled_topics | |
| for k in range(s_topics0.shape[-1]): | |
| topick0, topick1 = s_topics0[..., k], s_topics1[..., k] # [N, L+S] | |
| if (topick0.sum() > 0) and (topick1.sum() > 0): | |
| new_feat0, new_mask0, selected_ids0 = self.reduce_feat( | |
| feat0, topick0, N, C | |
| ) | |
| new_feat1, new_mask1, selected_ids1 = self.reduce_feat( | |
| feat1, topick1, N, C | |
| ) | |
| for idt in range(self.n_iter_topic_transformer): | |
| new_feat0 = self.topic_transformers[idt * 2]( | |
| new_feat0, new_feat0, new_mask0, new_mask0 | |
| ) | |
| new_feat1 = self.topic_transformers[idt * 2]( | |
| new_feat1, new_feat1, new_mask1, new_mask1 | |
| ) | |
| new_feat0 = self.topic_transformers[idt * 2 + 1]( | |
| new_feat0, new_feat1, new_mask0, new_mask1 | |
| ) | |
| new_feat1 = self.topic_transformers[idt * 2 + 1]( | |
| new_feat1, new_feat0, new_mask1, new_mask0 | |
| ) | |
| updated_feat0[selected_ids0, :] = new_feat0[new_mask0, :] | |
| updated_feat1[selected_ids1, :] = new_feat1[new_mask1, :] | |
| feat0 = (1 - s_topics0.sum(dim=-1, keepdim=True)) * feat0 + updated_feat0 | |
| feat1 = (1 - s_topics1.sum(dim=-1, keepdim=True)) * feat1 + updated_feat1 | |
| conf_matrix = ( | |
| torch.einsum("nlc,nsc->nls", feat0, feat1) / C**0.5 | |
| ) # (C * temperature) | |
| if self.training: | |
| topic_matrix = torch.einsum( | |
| "nlk,nsk->nls", prob_topics[:, :L], prob_topics[:, L:] | |
| ) | |
| outlier_mask = torch.einsum( | |
| "nlk,nsk->nls", feat_topics[:, :L], feat_topics[:, L:] | |
| ) | |
| else: | |
| topic_matrix = {"img0": feat_topics[:, :L], "img1": feat_topics[:, L:]} | |
| outlier_mask = torch.ones_like(conf_matrix) | |
| if mask0 is not None: | |
| outlier_mask = outlier_mask * mask0[..., None] * mask1[:, None] # .bool() | |
| conf_matrix.masked_fill_(~outlier_mask.bool(), -1e9) | |
| conf_matrix = F.softmax(conf_matrix, 1) * F.softmax( | |
| conf_matrix, 2 | |
| ) # * topic_matrix | |
| return feat0, feat1, conf_matrix, topic_matrix | |
| class LocalFeatureTransformer(nn.Module): | |
| """A Local Feature Transformer (LoFTR) module.""" | |
| def __init__(self, config): | |
| super(LocalFeatureTransformer, self).__init__() | |
| self.config = config | |
| self.d_model = config["d_model"] | |
| self.nhead = config["nhead"] | |
| self.layer_names = config["layer_names"] | |
| encoder_layer = LoFTREncoderLayer( | |
| config["d_model"], config["nhead"], config["attention"] | |
| ) | |
| self.layers = nn.ModuleList( | |
| [copy.deepcopy(encoder_layer) for _ in range(2)] | |
| ) # len(self.layer_names))]) | |
| self._reset_parameters() | |
| def _reset_parameters(self): | |
| for p in self.parameters(): | |
| if p.dim() > 1: | |
| nn.init.xavier_uniform_(p) | |
| def forward(self, feat0, feat1, mask0=None, mask1=None): | |
| """ | |
| Args: | |
| feat0 (torch.Tensor): [N, L, C] | |
| feat1 (torch.Tensor): [N, S, C] | |
| mask0 (torch.Tensor): [N, L] (optional) | |
| mask1 (torch.Tensor): [N, S] (optional) | |
| """ | |
| assert ( | |
| self.d_model == feat0.shape[2] | |
| ), "the feature number of src and transformer must be equal" | |
| feat0 = self.layers[0](feat0, feat1, mask0, mask1) | |
| feat1 = self.layers[1](feat1, feat0, mask1, mask0) | |
| return feat0, feat1 | |