4. 多模态Transformer架构实现

任务目标

实现一个能够同时处理图像与文本输入的多模态Transformer，输出两种模态融合后的联合表示。该表示可直接用于下游任务（如图文检索、视觉问答前的特征提取）。

核心实现

采用双流编码 + 交叉注意力融合的经典架构：图像经ViT风格patch嵌入，文本经嵌入，两者拼接后送入标准Transformer编码器进行模态间交互。以下代码为最小可复现实现，依赖PyTorch 2.0+和torchvision。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import numpy as np

# ---------- 1. 基础组件 ----------
class PatchEmbed(nn.Module):
    """图像分块嵌入: (B,C,H,W) -> (B,N,D)"""
    def __init__(self, img_size=224, patch_size=16, in_ch=3, embed_dim=256):
        super().__init__()
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_ch, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x)  # (B,D,14,14)
        x = x.flatten(2).transpose(1, 2)  # (B,N,D)
        return x

class TextEmbed(nn.Module):
    """文本嵌入 + 可学习位置编码"""
    def __init__(self, vocab_size=1000, max_len=32, embed_dim=256):
        super().__init__()
        self._emb = nn.Embedding(vocab_size, embed_dim)
        self.pos_emb = nn.Parameter(torch.zeros(1, max_len, embed_dim))
        nn.init.trunc_normal_(self.pos_emb, std=0.02)

    def forward(self, tokens):
        # tokens: (B,T) 整数索引
        x = self._emb(tokens)  # (B,T,D)
        x = x + self.pos_emb[:, :x.size(1), :]
        return x

class TransformerBlock(nn.Module):
    """标准Transformer编码层"""
    def __init__(self, dim=256, n_heads=4, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(int(dim * mlp_ratio), dim)
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = x + self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])
        x = x + self.dropout(self.mlp(self.norm2(x)))
        return x

# ---------- 2. 多模态Transformer ----------
class MultimodalTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3,
                 vocab_size=1000, text_len=32, embed_dim=256,
                 depth=4, n_heads=4, num_classes=10):
        super().__init__()
        self.patch_embed = PatchEmbed(img_size, patch_size, in_ch, embed_dim)
        self.text_embed = TextEmbed(vocab_size, text_len, embed_dim)
        self.cls_ = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.patch_embed.n_patches + text_len, embed_dim))
        nn.init.trunc_normal_(self.cls_, std=0.02)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, n_heads) for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)  # 可选分类头

    def forward(self, images, text_tokens):
        """
        images: (B,3,H,W) 浮点张量, 值范围[0,1]
        text_tokens: (B,T) 整数
        返回: (B, embed_dim) 融合表示
        """
        B = images.size(0)
        # 图像patch嵌入
        img_feat = self.patch_embed(images)  # (B,N_img,D)
        # 文本嵌入
        txt_feat = self.text_embed(text_tokens)  # (B,T,D)

        # 拼接CLS 
        cls_tokens = self.cls_.expand(B, -1, -1)
        x = torch.cat([cls_tokens, img_feat, txt_feat], dim=1)  # (B,1+N+T,D)
        x = x + self.pos_embed[:, :x.size(1), :]

        # 多层Transformer编码
        for block in self.blocks:
            x = block(x)

        x = self.norm(x)
        # 取CLS位置作为融合表示
        fused = x[:, 0, :]  # (B,D)
        return fused

# ---------- 3. 最小验证闭环 ----------
class DummyDataset(Dataset):
    """生成随机图像与文本用于验证前向传播"""
    def __init__(self, n_samples=4, img_size=224, text_len=8, vocab_size=1000):
        self.n = n_samples
        self.img_size = img_size
        self.text_len = text_len
        self.vocab_size = vocab_size

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        img = torch.rand(3, self.img_size, self.img_size)  # 随机图像
        tokens = torch.randint(0, self.vocab_size, (self.text_len,))
        return img, tokens

def verify():
    """运行前向传播并检查输出形状"""
    torch.manual_seed(42)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = MultimodalTransformer(
        img_size=224, patch_size=16, embed_dim=256,
        depth=4, n_heads=4, vocab_size=1000, text_len=8
    ).to(device)

    dataset = DummyDataset(n_samples=2)
    loader = DataLoader(dataset, batch_size=2)

    for images, tokens in loader:
        images, tokens = images.to(device), tokens.to(device)
        fused = model(images, tokens)
        print(f"输入图像: {tuple(images.shape)}, 文本: {tuple(tokens.shape)}")
        print(f"融合表示输出: {tuple(fused.shape)}")
        assert fused.shape == (2, 256), "输出形状应为 (batch, embed_dim)"
        print("✅ 前向传播验证通过")
        break

if __name__ == "__main__":
    verify()

运行与验证

1. 环境依赖：pip install torch torchvision pillow numpy
2. 运行方式：将上述代码保存为 mm_transformer.py，执行 python mm_transformer.py
3. 通过条件：控制台输出 融合表示输出: torch.Size([2, 256]) 且断言无异常，即证明模型能正确处理图像+文本输入并输出指定维度的融合特征。

关键说明

• 图像输入需为 (B,3,224,224) 浮点张量，值域[0,1]；文本为整数序列，长度≤32。
• 融合表示取自CLS 经过全部Transformer层后的输出，聚合了图像与文本的全局交互信息。
• 若要接入具体下游任务（如分类），可将 fused 输入到任务头；本案例仅验证融合表示的正确生成。
• 实际训练时需替换DummyDataset为真实图文对数据，并添加损失函数与优化器。