1. 多模态数据预处理与统一表示

任务目标

构建一个可复用的多模态数据预处理流水线,将图像、文本、音频三种模态的原始数据转换为统一的特征向量表示,并输出为可直接用于模型训练的批次数据。本案例聚焦预处理与统一表示环节,不涉及下游模型训练。

环境依赖

需要 torch、torchvision、transformers、librosa、Pillow。若使用CPU运行,需将模型参数设置为 device="cpu"。

核心实现

以下代码实现了一个多模态预处理类,分别通过预训练模型提取图像特征(ResNet18)、文本特征(BERT-base)和音频特征(Wav2Vec2),并将三者映射到同一维度空间(256维),最后通过一个简单的批处理函数将多个样本堆叠为统一张量。

import torch
import torch.nn as nn
from torchvision import models, transforms
from transformers import BertModel, BertTokenizer, Wav2Vec2Processor, Wav2Vec2Model
import librosa
import numpy as np
from PIL import Image

class UnifiedFeatureExtractor(nn.Module):
    def __init__(self, unified_dim=256, device="cpu"):
        super().__init__()
        self.device = device
        self.unified_dim = unified_dim

        # 图像分支:ResNet18 去掉全连接层
        resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.image_encoder = nn.Sequential(*list(resnet.children())[:-1])
        self.image_proj = nn.Linear(512, unified_dim)

        # 文本分支:BERT-base-uncased
        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
        bert = BertModel.from_pretrained("bert-base-uncased")
        self.text_encoder = bert
        self.text_proj = nn.Linear(768, unified_dim)

        # 音频分支:Wav2Vec2-base
        self.audio_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base")
        wav2vec = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base")
        self.audio_encoder = wav2vec
        self.audio_proj = nn.Linear(768, unified_dim)

        self.image_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        self.to(device)
        self.eval()

    @torch.no_grad()
    def encode_image(self, image_path):
        img = Image.open(image_path).convert("RGB")
        img = self.image_transform(img).unsqueeze(0).to(self.device)
        feat = self.image_encoder(img).flatten(1)
        return self.image_proj(feat)

    @torch.no_grad()
    def encode_text(self, text):
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True,
                                max_length=64, padding="max_length").to(self.device)
        outputs = self.text_encoder(**inputs)
        # 取 [CLS] 向量
        feat = outputs.last_hidden_state[:, 0, :]
        return self.text_proj(feat)

    @torch.no_grad()
    def encode_audio(self, audio_path, sample_rate=16000):
        speech, sr = librosa.load(audio_path, sr=sample_rate)
        inputs = self.audio_processor(speech, sampling_rate=sample_rate,
                                      return_tensors="pt", padding=True,
                                      truncation=True, max_length=16000).to(self.device)
        outputs = self.audio_encoder(**inputs)
        # 对时间维做均值池化
        feat = outputs.last_hidden_state.mean(dim=1)
        return self.audio_proj(feat)

    def forward(self, image_path=None, text=None, audio_path=None):
        """输入任意一个或多个模态路径,返回统一维度的特征字典"""
        result = {}
        if image_path is not None:
            result["image"] = self.encode_image(image_path)
        if text is not None:
            result["text"] = self.encode_text(text)
        if audio_path is not None:
            result["audio"] = self.encode_audio(audio_path)
        return result

def collate_modalities(sample_list):
    """将多个样本的模态特征堆叠为批次张量,缺失模态用零向量填充"""
    batch_size = len(sample_list)
    device = sample_list[0]["image"].device if "image" in sample_list[0] else torch.device("cpu")
    unified_dim = next(iter(sample_list[0].values())).shape[1]

    batch = {}
    for modality in ["image", "text", "audio"]:
        tensors = []
        for sample in sample_list:
            if modality in sample:
                tensors.append(sample[modality])
            else:
                tensors.append(torch.zeros(1, unified_dim, device=device))
        batch[modality] = torch.cat(tensors, dim=0)
    return batch

# 最小调用示例
if __name__ == "__main__":
    extractor = UnifiedFeatureExtractor(device="cpu")

    # 准备三个模态的示例输入(请替换为实际文件路径)
    sample1 = extractor(image_path="cat.jpg", text="a cat sitting on a chair", audio_path="speech.wav")
    sample2 = extractor(image_path="dog.jpg", text="a dog running in the park")

    # 组成批次
    batch = collate_modalities([sample1, sample2])
    print("Batch shapes:", {k: v.shape for k, v in batch.items()})
    # 预期输出: Batch shapes: {'image': torch.Size([2, 256]), 'text': torch.Size([2, 256]), 'audio': torch.Size([2, 256])}

运行与验证

运行方式:将上述代码保存为 multimodal_preprocess.py,在包含示例图片和音频文件的目录中执行 python multimodal_preprocess.py。首次运行会自动下载预训练模型权重。

验证方法:观察终端输出的批次张量形状。若三个模态的批次维度均为 (2, 256),说明流水线工作正常。其中第2个样本未提供音频,其音频特征被零向量填充,这验证了缺失模态的容错处理。

通过条件:程序无报错且输出的形状与预期一致,即所有模态特征均被映射到256维统一空间,且批处理函数能正确处理缺失模态。

关键说明

  • 三个模态的编码器均为冻结状态的预训练模型,仅投影层(Linear)参与训练,这保证了特征提取的稳定性并降低了计算开销。
  • 文本统一截断至64个,音频统一截断至1秒(16000采样点),超出部分被截断,不足部分以padding填充,确保单样本特征维度固定。
  • 实际训练时,只需将 extractor 的 eval() 改为 train(),并将投影层加入优化器即可端到端微调。