| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class V_Attention(nn.Module): |
| def __init__(self, dim_I, num_heads=8, dropout=0.1, ffn_ratio=4.0): |
| super(V_Attention, self).__init__() |
| self.dim_I = dim_I |
| self.num_heads = num_heads |
| self.dropout = nn.Dropout(dropout) |
| self.cross_attn_a_to_b = nn.MultiheadAttention( |
| embed_dim=dim_I, |
| num_heads=num_heads, |
| kdim=dim_I, |
| vdim=dim_I, |
| dropout=dropout, |
| batch_first=True |
| ) |
| self.ffn_i = nn.Sequential( |
| nn.Linear(dim_I, dim_I*4), |
| nn.GELU(), |
| nn.Dropout(0.1), |
| nn.Linear(dim_I*4, dim_I), |
| nn.LayerNorm(dim_I), |
| ) |
| self.norm1_a = nn.LayerNorm(dim_I) |
|
|
| def forward(self, feat_I, mask=None): |
| B, N_a, _ = feat_I.shape |
| out_a, attn_a = self.cross_attn_a_to_b( |
| query=feat_I, |
| key=feat_I, |
| value=feat_I, |
| key_padding_mask=~mask, |
| need_weights=True |
| ) |
| out_a = feat_I + self.dropout(self.ffn_i(out_a)) |
| out_I = self.norm1_a(out_a) |
| return out_I |
|
|
|
|
| class VTS_Alignment(nn.Module): |
| def __init__(self, v_dim, TS_dim, embedding_dim=512, num_heads=8, dropout=0.1): |
| super().__init__() |
| self.TS_dim = TS_dim |
| self.embedding_dim = embedding_dim |
|
|
| t_dim2 = int(2 * self.TS_dim) |
| self.mlp_i = nn.Sequential( |
| nn.Linear(TS_dim, t_dim2), |
| nn.GELU(), |
| nn.Linear(t_dim2, TS_dim), |
| nn.LayerNorm(TS_dim), |
| ) |
|
|
| self.mlp_t = nn.Sequential( |
| nn.Linear(TS_dim, t_dim2), |
| nn.GELU(), |
| nn.Linear(t_dim2, TS_dim), |
| nn.LayerNorm(TS_dim), |
| ) |
|
|
| self.cross_attn_a_to_b = nn.MultiheadAttention( |
| embed_dim=TS_dim, |
| num_heads=num_heads, |
| kdim=TS_dim, |
| vdim=TS_dim, |
| dropout=dropout, |
| batch_first=True |
| ) |
|
|
| self.cross_attn_b_to_a = nn.MultiheadAttention( |
| embed_dim=TS_dim, |
| num_heads=num_heads, |
| kdim=TS_dim, |
| vdim=TS_dim, |
| dropout=dropout, |
| batch_first=True |
| ) |
|
|
| self.ffn_a = nn.Sequential( |
| nn.Linear(TS_dim, t_dim2), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(t_dim2, TS_dim) |
| ) |
|
|
| self.ffn_b = nn.Sequential( |
| nn.Linear(TS_dim, t_dim2), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(t_dim2, TS_dim) |
| ) |
|
|
| self.norm1_a = nn.LayerNorm(TS_dim) |
| self.norm2_a = nn.LayerNorm(TS_dim) |
| self.norm1_b = nn.LayerNorm(TS_dim) |
| self.norm2_b = nn.LayerNorm(TS_dim) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, feat_I, feat_TS, mask=None): |
| feat_I = self.mlp_i(feat_I) |
| feat_TS = self.mlp_t(feat_TS) |
| |
| out_a, _ = self.cross_attn_a_to_b( |
| query=feat_I, |
| key=feat_TS, |
| value=feat_TS, |
| key_padding_mask=~mask, |
| need_weights=True |
| ) |
| out_a = feat_I + self.dropout(out_a) |
| out_a = self.norm1_a(out_a) |
| out_a = out_a + self.dropout(self.ffn_a(out_a)) |
| x_I = self.norm2_a(out_a) |
|
|
| out_b, attn_b = self.cross_attn_b_to_a( |
| query=feat_TS, |
| key=feat_I, |
| value=feat_I, |
| key_padding_mask=~mask, |
| need_weights=True |
| ) |
| out_b = feat_TS + self.dropout(out_b) |
| out_b = self.norm1_b(out_b) |
| out_b = out_b + self.dropout(self.ffn_b(out_b)) |
| x_TS = self.norm2_b(out_b) |
|
|
| return x_I, x_TS |
|
|
|
|
| class router(nn.Module): |
| def __init__(self, dim, channel_num, num_tasks=2, topk=2, task_model='complex'): |
| super().__init__() |
| embed_dim = int(dim // 8) |
| self.task_model = task_model |
| self.l1 = nn.Linear(dim, embed_dim) |
| self.l2 = nn.Linear(embed_dim, channel_num) |
| self.topk = topk |
| self.task_embedding = nn.Embedding(num_tasks, embed_dim) |
|
|
| def forward(self, x, task_id=None): |
| original_shape = x.shape |
| x = x.view(-1, x.size(-1)) |
| x = F.gelu(self.l1(x)) |
| |
| if task_id is not None: |
| task_id = int(task_id) |
| task_id = torch.tensor(task_id, device=x.device, dtype=torch.long) |
| task_bias = self.task_embedding(task_id) |
| task_emb = task_bias.unsqueeze(0).expand(x.size(0), -1) |
| x = x + task_emb |
|
|
| logits = self.l2(x) |
| topk_vals, topk_idx = torch.topk(logits, self.topk, dim=-1) |
| topk_probs = torch.softmax(topk_vals, dim=-1).to(logits.dtype) |
| probs = torch.zeros_like(logits) |
| probs.scatter_(-1, topk_idx, topk_probs) |
| probs = probs.view(*original_shape[:-1], -1) |
| return probs |
|
|
|
|
| class M_moe(nn.Module): |
| def __init__(self, dst_feature_dims): |
| super(M_moe, self).__init__() |
| self.dims = dst_feature_dims |
| self.Router = router(self.dims * 2, 3) |
| self.mlp_m = nn.Sequential( |
| nn.Linear(self.dims * 2, self.dims * 2), |
| nn.GELU(), |
| nn.Dropout(0.1), |
| nn.Linear(self.dims * 2, self.dims), |
| ) |
|
|
| def forward(self, F_M, F_T, F_I, router_input, mask=None): |
| B, T, _ = F_T.shape |
| F_M = self.mlp_m(F_M) |
| if mask is not None: |
| m_w_r = self.Router(router_input, 0) |
| c_fusion = ( |
| F_T * m_w_r[:, :, 0:1] + |
| F_M * m_w_r[:, :, 1:2] + |
| F_I * m_w_r[:, :, 2:3] |
| ) |
| return c_fusion, m_w_r |
| else: |
| m_w_c = self.Router(router_input, 1) |
| c_fusion = ( |
| F_T * m_w_c[:, :, 0:1] + |
| F_M * m_w_c[:, :, 1:2] + |
| F_I * m_w_c[:, :, 2:3] |
| ) |
| return c_fusion, m_w_c |
|
|
