File size: 1,311 Bytes
12d60da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | import torch
import torch.nn as nn
class CrossAttentionFusion(nn.Module):
"""
Fuses features from two different sources (e.g. CNN features and ViT/Transformer features)
using Cross-Attention mechanism.
"""
def __init__(self, d_model=1280, nhead=8, dropout=0.1):
super().__init__()
self.multihead_attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout, batch_first=True)
self.norm = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, query, key, value):
# query, key, value shape: [batch_size, seq_len, d_model]
attn_output, _ = self.multihead_attn(query, key, value)
x = query + self.dropout(attn_output)
x = self.norm(x)
return x
class SimpleConcatFusion(nn.Module):
"""
Simple concatenation of features followed by linear projection.
"""
def __init__(self, in_features1, in_features2, out_features):
super().__init__()
self.fc = nn.Linear(in_features1 + in_features2, out_features)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.3)
def forward(self, feat1, feat2):
x = torch.cat([feat1, feat2], dim=-1)
x = self.fc(x)
x = self.relu(x)
x = self.dropout(x)
return x
|