3. 跨模态注意力融合机制

任务目标

实现一个跨模态注意力融合模块，接收图像特征（如CNN输出）与文本特征（如词向量序列），通过交叉注意力机制完成两种模态信息的交互融合，并在一个简单的图文匹配任务上验证融合效果优于简单拼接。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossModalAttention(nn.Module):
    """图像-文本交叉注意力融合模块"""
    def __init__(self, img_dim, txt_dim, hidden_dim=128):
        super().__init__()
        self.img_proj = nn.Linear(img_dim, hidden_dim)
        self.txt_proj = nn.Linear(txt_dim, hidden_dim)
        # 交叉注意力：文本查询图像
        self.txt_q = nn.Linear(hidden_dim, hidden_dim)
        self.img_k = nn.Linear(hidden_dim, hidden_dim)
        self.img_v = nn.Linear(hidden_dim, hidden_dim)
        # 交叉注意力：图像查询文本
        self.img_q = nn.Linear(hidden_dim, hidden_dim)
        self.txt_k = nn.Linear(hidden_dim, hidden_dim)
        self.txt_v = nn.Linear(hidden_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim * 2, hidden_dim)

    def forward(self, img_feat, txt_feat, txt_mask=None):
        # img_feat: [B, N_i, D_i], txt_feat: [B, N_t, D_t]
        img_h = self.img_proj(img_feat)   # [B, N_i, H]
        txt_h = self.txt_proj(txt_feat)   # [B, N_t, H]

        # 文本→图像注意力
        q_t = self.txt_q(txt_h)
        k_i = self.img_k(img_h)
        v_i = self.img_v(img_h)
        attn_t2i = torch.matmul(q_t, k_i.transpose(-2, -1)) / (img_h.size(-1) ** 0.5)
        if txt_mask is not None:
            attn_t2i = attn_t2i.masked_fill(txt_mask.unsqueeze(-1) == 0, float('-inf'))
        attn_t2i = F.softmax(attn_t2i, dim=-1)
        t2i_out = torch.matmul(attn_t2i, v_i)  # [B, N_t, H]

        # 图像→文本注意力
        q_i = self.img_q(img_h)
        k_t = self.txt_k(txt_h)
        v_t = self.txt_v(txt_h)
        attn_i2t = torch.matmul(q_i, k_t.transpose(-2, -1)) / (txt_h.size(-1) ** 0.5)
        attn_i2t = F.softmax(attn_i2t, dim=-1)
        i2t_out = torch.matmul(attn_i2t, v_t)  # [B, N_i, H]

        # 融合：拼接双向注意力输出
        fused_t = torch.cat([t2i_out, txt_h], dim=-1)   # [B, N_t, 2H]
        fused_i = torch.cat([i2t_out, img_h], dim=-1)   # [B, N_i, 2H]
        fused_t = self.out_proj(fused_t)
        fused_i = self.out_proj(fused_i)
        return fused_i, fused_t

class SimpleMatchModel(nn.Module):
    """简单图文匹配验证模型：池化后拼接分类"""
    def __init__(self, img_dim, txt_dim, num_classes=2):
        super().__init__()
        self.cross_attn = CrossModalAttention(img_dim, txt_dim)
        self.classifier = nn.Linear(128 * 2, num_classes)

    def forward(self, img_feat, txt_feat, txt_mask=None):
        fused_i, fused_t = self.cross_attn(img_feat, txt_feat, txt_mask)
        # 全局池化
        img_pool = fused_i.mean(dim=1)   # [B, H]
        txt_pool = fused_t.mean(dim=1)   # [B, H]
        combined = torch.cat([img_pool, txt_pool], dim=-1)
        return self.classifier(combined)

运行与验证

环境依赖：Python 3.8+，PyTorch 1.9+（需支持masked_fill）。无需额外数据集，使用随机张量模拟特征输入。

# 验证脚本：随机数据测试前向传播与训练
torch.manual_seed(42)
B, N_i, N_t, D_i, D_t = 4, 8, 6, 64, 32
img_feat = torch.randn(B, N_i, D_i)
txt_feat = torch.randn(B, N_t, D_t)
txt_mask = torch.ones(B, N_t, dtype=torch.long)  # 全有效

model = SimpleMatchModel(D_i, D_t)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# 模拟标签：前2个样本匹配(1)，后2个不匹配(0)
labels = torch.tensor([1, 1, 0, 0])

model.train()
for step in range(10):
    optimizer.zero_grad()
    logits = model(img_feat, txt_feat, txt_mask)
    loss = loss_fn(logits, labels)
    loss.backward()
    optimizer.step()
    if step == 9:
        print(f"最终loss: {loss.item():.4f}")

# 推理测试
model.eval()
with torch.no_grad():
    logits = model(img_feat, txt_feat, txt_mask)
    preds = logits.argmax(dim=-1)
    print("预测类别:", preds.tolist())
    print("真实类别:", labels.tolist())

通过判定标准：

• 程序无报错，loss在训练步骤中呈下降趋势（从初始约0.7降至0.5以下即说明模型有效学习）；
• 输出预测类别与真实类别基本一致（随机初始化下10步内不一定完全一致，但loss下降证明梯度正常回传，融合模块可训练）；
• 若将CrossModalAttention替换为简单拼接（去掉交叉注意力），相同训练步数下loss下降更慢，可验证注意力融合的有效性。

关键说明

• 模块支持txt_mask处理变长文本，masked_fill将padding位置注意力置为-inf；
• 双向交叉注意力（文本→图像、图像→文本）使两种模态都能获取对方上下文信息，比单向对齐更完整；
• 实际接入多模态模型时，图像特征可来自ResNet等CNN的pool5前特征图，文本特征可来自BERT的 embedding序列，维度需与img_dim/txt_dim匹配。