多模态大模型MLLM基础训练范式 Pre-train + Instruction FineTuning

多模态大模型Pre-train

为了在图文嵌入空间中更好地对齐视觉和文本信息。为此，使用图像-文本对（image-caption style data），表示为 $(\mathbf{X},Y_a)$，其中：

• $\mathbf{X}$：图像（Image）。
• Y a = { y i } i = 1 N a Y_a = \{y_i\}_{i=1}^{N_a} Ya​={yi​}i=1Na​​：是图像的描述（response，描述或文本回应），包含 $N_a$ 个token（词元）。

---

下面两个公式是图文对齐预训练中核心的 语言建模概率函数 和 损失函数定义。它们共同描述了：如何让模型学会根据一张图像生成对应的文字描述，以实现图文信息在嵌入空间的对齐。

公式（1）：生成概率公式

$$p(Y_a|\mathbf{X})=\prod _{i=1}^{N_a}{F}_{\theta }(y_i|P_{\varphi }\circ V_{\phi }(\mathbf{X}))\text{(1)}$$

🔍 解释：

这是一个 条件语言建模公式（Conditional Language Modeling）。该公式定义了生成文本 $Y_a$ 在图像 $\mathbf{X}$ 条件下的概率，是一个自回归语言建模形式（Autoregressive Language Modeling）：

• ${\prod }_{i=1}^{N_a}$：表示对所有token做乘积，相当于整体生成序列的联合概率。
• $F_{\theta }$：是语言模型（如 LLM）中的解码器函数，参数为 $\theta$。
• $y_i$：目标文本序列的第 $i$ 个token。
• $P_{\varphi }$：是将视觉特征投影到语言模型空间的投影器（如一个MLP或线性变换）。
• $V_{\phi }(\mathbf{X})$：是视觉编码器（Vision Encoder），将图像 $\mathbf{X}$ 编码为视觉特征。
• $\circ$：表示函数复合，$P_{\varphi }\circ V_{\phi }(\mathbf{X})$ 就是“先编码图像再投影”。

---

• 给定图像 $\mathbf{X}$，目标是根据图像生成一段文字 Y a = { y 1 , y 2 , … , y N a } Y_a = \{y_1, y_2, \dots, y_{N_a}\} Ya​={y1​,y2​,…,yNa​​}。
• 每个词 $y_i$ 的生成概率，依赖于图像特征经过视觉编码器 $V_{\phi }$ 和投影模块 $P_{\varphi }$ 得到的特征。
• $F_{\theta }$ 是 LLM 的解码器部分，根据图像特征一步步生成句子。

---

公式（2）：训练目标（最大似然）

这个是对应的训练目标——最大化似然（Maximum Likelihood Estimation, MLE）。通过这种损失训练，视觉特征（来自图像）和语言特征（来自文本）被映射到同一个语义空间中，实现图文特征的对齐。
$$\underset{\varphi ,{\theta }^{\prime },{\phi }^{\prime }}{\max }\sum _{i=1}^{N_a}\log F_{\theta }(y_i|P_{\varphi }\circ V_{\phi }(\mathbf{X}))\text{(2)}$$

🔍 解释：

• $\log F_{\theta }(\cdot )$：表示的是每个token的对数概率（log-likelihood）。
• 对所有token求和即为整个序列的log-likelihood。
• ${\max }_{\varphi ,{\theta }^{\prime },{\phi }^{\prime }}$：表示对这三个部分的参数进行优化：
  • $\varphi$：投影器的参数。
  • ${\theta }^{\prime }$：语言模型中允许学习的部分参数（不是全部）。
  • ${\phi }^{\prime }$：视觉编码器中允许学习的部分参数（不是全部）。

损失函数定义（Loss Function）：

损失函数是上面最大化目标的对立面——我们在训练中实际上是最小化负对数似然（Negative Log-Likelihood Loss，NLL）：

$${\mathcal{L}}_{\text{NLL}}=-\sum _{i=1}^{N_a}\log F_{\theta }(y_i|P_{\varphi }\circ V_{\phi }(\mathbf{X}))$$

这个损失函数的核心思想就是：

• 如果模型对某个正确token的预测概率高（接近1），它的log值就是负的很小；
• 如果模型预测错了，log值就是负的很大，总损失就会变大；
• 所以，通过最小化负log概率，我们可以让模型学会尽量准确地生成目标描述文本。

---

• 用的是自回归语言模型的训练方式：一个词一个词预测。
• 每预测一个token $y_i$，我们就计算它的对数概率 $\log F_{\theta }(y_i|\cdot )$。
• 然后将所有token的log概率加起来。
• 最后最大化这个和（就是最小化负的log-likelihood）。目标是让正确的文本描述在图像条件下的生成概率越高越好。

⚠️ 说明：

• 只对部分参数进行训练，是因为使用的是小型语言模型（small-scale LLM），为了效率和防止过拟合，作者只更新关键连接器部分参数。这被称为Partial Parameter Tuning（部分参数可学习）。
• 仅仅训练连接器（如 $P_{\varphi }$）可能不足以实现图文对齐，因此本方法允许在小模型上微调部分 LLM 和视觉模型的参数。

---

Pre-train案例代码

import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2LMHeadModel, GPT2Tokenizer, ViTModel, ViTImageProcessor
import requests
import zipfile
from tqdm import tqdm
from PIL import Image
import json# ==================== 配置参数 ====================
class Config:# 数据配置coco_path = "./coco"image_size = 224       # ViT标准输入尺寸batch_size = 8         # 根据GPU显存调整max_seq_len = 64       # 文本序列最大长度# 模型配置vision_model = "google/vit-base-patch16-224-in21k"language_model = "gpt2"projection_dim = 768   # 投影层维度需与语言模型匹配# 训练配置lr = 1e-4epochs = 2           # 训练轮数,根据数据集大小和GPU显存调整device = torch.device("cuda" if torch.cuda.is_available() else "cpu")grad_clip = 1.0        # 梯度裁剪阈值num_workers = 4        # 数据加载线程数# ==================== 数据下载 ====================
def download_with_progress(url: str, save_path: str):"""带进度条的文件下载函数"""response = requests.get(url, stream=True)total_size = int(response.headers.get('content-length', 0))with open(save_path, 'wb') as f, tqdm(desc=f"Downloading {os.path.basename(save_path)}",total=total_size,unit='iB',unit_scale=True,unit_divisor=1024,) as bar:for data in response.iter_content(chunk_size=1024):f.write(data)bar.update(len(data))def prepare_coco_data():"""准备COCO数据集"""os.makedirs(Config.coco_path, exist_ok=True)# 下载训练图像image_zip = os.path.join(Config.coco_path, "train2017.zip")if not os.path.exists(os.path.join(Config.coco_path, "train2017")):print("下载COCO训练集...")download_with_progress("http://images.cocodataset.org/zips/train2017.zip",image_zip)with zipfile.ZipFile(image_zip, 'r') as zip_ref:zip_ref.extractall(Config.coco_path)os.remove(image_zip)# 下载标注文件ann_zip = os.path.join(Config.coco_path, "annotations.zip")if not os.path.exists(os.path.join(Config.coco_path, "annotations")):print("下载COCO标注...")download_with_progress("http://images.cocodataset.org/annotations/annotations_trainval2017.zip",ann_zip)with zipfile.ZipFile(ann_zip, 'r') as zip_ref:zip_ref.extractall(Config.coco_path)os.remove(ann_zip)# ==================== 数据集类 ====================
class CocoCaptionDataset(Dataset):"""COCO图像描述数据集"""def __init__(self, image_dir: str, annotation_path: str):self.image_dir = image_dir# 初始化图像处理器try:self.image_processor = ViTImageProcessor.from_pretrained(Config.vision_model)except Exception as e:raise RuntimeError(f"无法加载图像处理器: {str(e)}")# 加载标注数据try:with open(annotation_path) as f:annotations = json.load(f)except FileNotFoundError:raise FileNotFoundError(f"标注文件 {annotation_path} 不存在")# 构建数据映射self.id_to_filename = {img["id"]: img["file_name"] for img in annotations["images"]}self.annotations = [(ann["image_id"], ann["caption"]) for ann in annotations["annotations"]]# 初始化文本处理器try:self.tokenizer = GPT2Tokenizer.from_pretrained(Config.language_model)# 使用EOS作为填充符，确保模型正确处理填充if self.tokenizer.pad_token is None:self.tokenizer.pad_token = self.tokenizer.eos_tokenexcept Exception as e:raise RuntimeError(f"无法加载文本处理器: {str(e)}")def __len__(self):return len(self.annotations)def __getitem__(self, idx):image_id, caption = self.annotations[idx]# 加载并处理图像image_path = os.path.join(self.image_dir, self.id_to_filename[image_id])try:image = Image.open(image_path).convert("RGB")pixel_values = self.image_processor(images=image, return_tensors="pt").pixel_values.squeeze(0)  # (3, 224, 224)except Exception as e:raise RuntimeError(f"图像处理失败: {image_path} - {str(e)}")# 处理文本tokens = self.tokenizer(caption,max_length=Config.max_seq_len,padding="max_length",truncation=True,return_tensors="pt")return pixel_values, tokens.input_ids.squeeze(0), tokens.attention_mask.squeeze(0)# ==================== 多模态模型 ====================
class VisionLanguageModel(nn.Module):"""视觉-语言多模态模型"""def __init__(self):super().__init__()# 视觉编码器（冻结）try:self.vision_encoder = ViTModel.from_pretrained(Config.vision_model)for param in self.vision_encoder.parameters():param.requires_grad = Falseexcept Exception as e:raise RuntimeError(f"无法加载视觉模型: {str(e)}")# 视觉特征投影层self.visual_projection = nn.Sequential(nn.Linear(self.vision_encoder.config.hidden_size, Config.projection_dim),nn.GELU(),nn.LayerNorm(Config.projection_dim))# 语言模型（部分微调）try:self.language_model = GPT2LMHeadModel.from_pretrained(Config.language_model)# 冻结除最后两层外的所有参数for param in self.language_model.parameters():param.requires_grad = Falsefor param in self.language_model.transformer.h[-2:].parameters():param.requires_grad = True# 获取GPT2分词器以保持一致性self.tokenizer = GPT2Tokenizer.from_pretrained(Config.language_model)if self.tokenizer.pad_token is None:self.tokenizer.pad_token = self.tokenizer.eos_tokenself.language_model.config.pad_token_id = self.language_model.config.eos_token_idexcept Exception as e:raise RuntimeError(f"无法加载语言模型: {str(e)}")def forward(self, images, input_ids, attention_mask):batch_size = images.size(0)# 视觉特征提取vision_outputs = self.vision_encoder(images)visual_features = vision_outputs.last_hidden_state[:, 0, :]  # [CLS] token# 特征投影projected_visual = self.visual_projection(visual_features).unsqueeze(1)  # [B, 1, dim]# 修改标签以匹配维度# 将输入的后移一位作为标签，并忽略视觉标记的位置labels = input_ids.clone()  # 复制输入作为标签# 将标签中所有视觉标记的位置设为-100，使损失函数忽略这些位置labels_attention_mask = torch.cat([torch.zeros(batch_size, 1, device=Config.device),  # 视觉标记部分attention_mask  # 文本部分], dim=1)# 文本嵌入text_embeddings = self.language_model.transformer.wte(input_ids)# 合并视觉和文本特征combined_embeddings = torch.cat([projected_visual, text_embeddings], dim=1)# 调整注意力掩码visual_attention = torch.ones(batch_size, 1, device=Config.device)combined_attention = torch.cat([visual_attention, attention_mask], dim=1)# 使用自定义标签进行损失计算outputs = self.language_model(inputs_embeds=combined_embeddings,attention_mask=combined_attention,labels=None  # 不使用内部标签处理)# 手动计算损失logits = outputs.logits  # [B, seq_len+1, vocab_size]# 移除视觉标记的logits，只保留文本部分shift_logits = logits[:, :-1, :]  # [B, seq_len, vocab_size]# 为了计算损失，我们将输入向右移动一位，这样每个token都是预测下一个token的shift_labels = torch.cat([input_ids[:, 1:],  # 将输入向右移动一位torch.full((batch_size, 1), self.tokenizer.eos_token_id, device=Config.device)  # 添加EOS作为最后一个标签], dim=1)# 忽略填充标记的损失loss_mask = attention_mask.float()# 计算交叉熵损失loss_fct = nn.CrossEntropyLoss(reduction='none')loss = loss_fct(shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1))# 应用掩码并计算平均值masked_loss = (loss.view(batch_size, -1) * loss_mask).sum() / loss_mask.sum()return masked_loss# ==================== 训练流程 ====================
def train():# 准备数据prepare_coco_data()# 初始化数据集try:train_dataset = CocoCaptionDataset(image_dir=os.path.join(Config.coco_path, "train2017"),annotation_path=os.path.join(Config.coco_path, "annotations/captions_train2017.json"))except Exception as e:print(f"数据集初始化失败: {str(e)}")return# 数据加载器train_loader = DataLoader(train_dataset,batch_size=Config.batch_size,shuffle=True,num_workers=Config.num_workers,pin_memory=True)# 初始化模型try:model = VisionLanguageModel().to(Config.device)except Exception as e:print(f"模型初始化失败: {str(e)}")return# 优化器配置optimizer = torch.optim.AdamW([{'params': model.visual_projection.parameters()},{'params': model.language_model.transformer.h[-2:].parameters()}], lr=Config.lr)# 学习率调度scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Config.epochs)# 混合精度训练# scaler = torch.amp.GradScaler()scaler = torch.cuda.amp.GradScaler()  # 启用混合精度训练# 训练循环for epoch in range(Config.epochs):model.train()total_loss = 0.0progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{Config.epochs}")for batch_idx, (images, input_ids, attention_mask) in enumerate(progress_bar):# 数据转移到设备images = images.to(Config.device, non_blocking=True)input_ids = input_ids.to(Config.device, non_blocking=True)attention_mask = attention_mask.to(Config.device, non_blocking=True)# 梯度清零optimizer.zero_grad(set_to_none=True)# 混合精度前向with torch.amp.autocast(device_type='cuda', dtype=torch.float16):loss = model(images, input_ids, attention_mask)# 反向传播scaler.scale(loss).backward()# 梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), Config.grad_clip)# 参数更新scaler.step(optimizer)scaler.update()# 统计损失total_loss += loss.item()progress_bar.set_postfix({"batch_loss": loss.item()})# # 添加定期检查点（每1000个批次保存一次）# if (batch_idx + 1) % 1000 == 0:#     torch.save({#         'epoch': epoch,#         'iteration': batch_idx,#         'model_state_dict': model.state_dict(),#         'optimizer_state_dict': optimizer.state_dict(),#         'loss': loss.item(),#     }, f"coco_caption_epoch_{epoch}_iter_{batch_idx}.pt")# 更新学习率scheduler.step()# 保存检查点torch.save({'epoch': epoch,'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss': total_loss / len(train_loader),}, f"coco_caption_epoch_{epoch}.pt")# 打印统计信息avg_loss = total_loss / len(train_loader)print(f"\nEpoch {epoch+1} | Avg Loss: {avg_loss:.4f} | last_LR: {scheduler.get_last_lr()[0]:.2e}")if __name__ == "__main__":train()

---

多模态大模型Instruction FineTuning

监督式微调。我们使用图像-文本对 $(X,Y)$ 作为多轮对话的原始形式。令 $\mathcal{A}$表示属于助手回应的所有标记的集合，即 A = { y ∣ y ∈ Y t a , ∀ t = 1 , … , T } \mathcal{A} = \{y \mid y \in Y_t^a, \, \forall t =1, \dots, T\} A={y∣y∈Yta​,∀t=1,…,T}，其中 $Y_t^a$ 是第 $t$轮对话中的助手回应。我们通过最大化助手回应的对数似然，逐步地将其作为训练目标进行监督式微调(最大化模型回应的对数似然)
$$\underset{\varphi ,{\theta }^{\prime },\phi }{\max }\sum _{i=1}^N\mathbb{I}(y_i\in \mathcal{A})\log F_{\theta }(y_i|P_{\varphi }\circ V_{\phi }(X))$$其中 $N$ 是文本序列 $Y$的长度，$\mathbb{I}(y_i\in \mathcal{A})$ 是指示函数，当 $y_i\in \mathcal{A}$ 时为1，否则为 0。我们还允许在监督式微调阶段对语言模型和视觉编码器的部分可学习参数进行调整。

---

公式（1）有选择的语言建模目标

$$\underset{\varphi ,\theta ,\phi }{\max }\sum _{i=1}^N\mathbb{I}(y_i\in \mathcal{A})\log F_{\theta }(y_i|P_{\varphi }\circ V_{\phi }(X))\text{(1)}$$

其中包含以下几个重要部分：

• ${\max }_{\varphi ,\theta ,\phi }$：表示我们正在优化（最大化）目标函数的参数，具体来说，优化的是模型的三个参数：$\varphi$（连接器参数）、$\theta$（语言模型参数）和 $\phi$（视觉编码器参数）。
• ${\sum }_{i=1}^N$：对整个输出文本序列 $Y$（长度为 $N$）的每一个元素进行求和，其中 $N$ 是文本序列的长度。
• $\mathbb{I}(y_i\in \mathcal{A})$：这是指示函数（indicator function），当 $y_i\in \mathcal{A}$ 时，其值为 1，否则为 0。$\mathcal{A}$ 表示属于助手回应的标记集合，用于确认每个 $y_i$ 是否是来自助手的回应。
• $F_{\theta }(y_i|P_{\varphi }\circ V_{\phi }(X))$：这是生成模型的输出，表示给定图像 $X$，通过视觉编码器 $V_{\phi }(X)$ 和连接器 $P_{\varphi }$，由语言模型 $F_{\theta }$ 生成文本 $y_i$ 的概率。这个概率依赖于图像 $X$ 以及模型参数 $\theta$、$\varphi$ 和 $\phi$。
  • $F_{\theta }$：小型语言模型（LLM），其参数为 $\theta$。
  • $P_{\varphi }$：连接器，用于连接视觉编码器和语言模型，其参数为 $\varphi$。
  • $V_{\phi }$：视觉编码器，用于处理图像数据，其参数为 $\phi$。
  • $\circ$：表示视觉编码器 $V_{\phi }$ 和连接器 $P_{\varphi }$ 之间的组合操作。

---

这是一个 有选择的语言建模目标（Selective Language Modeling Objective）
和普通的语言建模一样，我们最大化目标 token 的对数概率（log-likelihood），但它有两点关键不同：

1. 选择性训练（Selective Training）：

   • 并不是对整个文本序列都训练，而是只训练那些属于助手回应的部分$\mathcal{A}$的 token。
   • 通过 $\mathbb{I}(y_i\in \mathcal{A})$这个指示函数实现。

2. 多模态条件语言建模（Conditioned on Image）：

   • 文本生成不仅依赖前面的文本，还要依赖图像 $X$ 的表示。
   • 图像通过视觉编码器 $V_{\phi }$ 和连接器 $P_{\varphi }$ 映射为适合语言模型输入形式。

条件语言建模

$F_{\theta }(y_i\mid y_{<i},Z)$这是语言模型预测 token $y_i$ 的概率，前提是：

• 已知图像表示 $Z$
• 已知之前已生成的 token 序列 $y_{<i}$
• 实际上，它可以表示为：$\log P(y_1,y_2,\dots ,y_N\mid Z)={\sum }_{i=1}^N\log P(y_i\mid y_{<i},Z)$

加入掩码（只训练助手回应的 token）

我们不对所有 token 求损失，而是只计算：$y_i\in \mathcal{A}$，即属于助手回应的那部分的 token。这通过乘以一个指示函数实现：
$$\mathbb{I}(y_i\in \mathcal{A})=\{\begin{array}{ll}1& \text{如果 }{y}_i\text{ 是助手回应的 token}\\ 0& \text{否则}\end{array}$$

公式（2）损失函数

最终的训练目标是：

$$\mathcal{L}=-\sum _{i=1}^N\mathbb{I}(y_i\in \mathcal{A})\cdot \log P(y_i\mid y_{<i},Z)\text{(2)}$$

• 模型需要在给定图像表示 $Z$ 和之前生成的 token 序列 $y_{<i}$ 的条件下，生成下一个 token $y_i$。因此，我们采用负对数似然损失（negative log-likelihood）来衡量生成正确 token 的概率，即 $-\log P(y_i\mid y_{<i},Z)$。
• 在训练对话模型时，输入通常包含了用户的提示和系统的回答。我们只希望模型针对“助手回应”的部分（记作集合 $\mathcal{A}$）进行学习，而不希望对那些已存在的提示部分产生损失。通过使用指示函数 $\mathbb{I}(y_i\in \mathcal{A})$，我们只对回答部分的 token 计算损失，从而防止模型“记住”提示内容，而专注于生成高质量的回复。
• 整个训练过程的核心是让模型学会如何根据图像信息 $Z$ 和上下文生成恰当的回答。通过仅在助手生成部分计算损失，我们确保训练信号只来自模型需要生成的内容，从而更好地实现视觉与语言之间的对齐和有效指令跟随。

分析指令微调公式（1）和公式（2）的联系和区别

这两个公式从本质上都在优化多模态大模型在“助手回答”部分生成正确文本的概率，通过对这些 token 的对数概率求和，可以鼓励模型输出高概率的（即正确的）回答。但它们表达方式和细节略有不同。

公式1：
$$\underset{\varphi ,\theta ,\phi }{\max }\sum _{i=1}^N\mathbb{I}(y_i\in \mathcal{A})\log F_{\theta }(y_i\mid P_{\varphi }\circ V_{\phi }(X))$$
这个公式以最大化形式出现，直接表达了“最大化对数似然”的目标。其中：

• $V_{\phi }(X)$ 表示将图像 $X$ 编码为视觉特征；
• $P_{\varphi }$ 将视觉特征投影到语言模型的嵌入空间；
• $F_{\theta }$ 则是整个语言生成模型，它隐式地包含了自回归结构（即内部考虑了上下文 $y_{<i}$），因此公式中没有显式写出 $y_{<i}$。

公式2：
$$\mathcal{L}=-\sum _{i=1}^N\mathbb{I}(y_i\in \mathcal{A})\cdot \log P(y_i\mid y_{<i},Z)$$
这个公式以损失函数的形式给出，采用了负对数似然（negative log-likelihood, NLL）的定义，目标是最小化这个损失。这里：

• $Z$ 通常代表经过投影后的视觉特征，即 $Z=P_{\varphi }\circ V_{\phi }(X)$；
• 显式地包含了自回归条件 $y_{<i}$，说明在生成每个 token 时，模型会考虑之前已经生成的 token。

优化目标

• 公式1是一个最大化问题，直接追求对数似然的最大化。
• 公式2则是一个最小化问题，即最小化负对数似然损失。从数学上讲，这两种表述在目标上是等价的（最大化对数似然等价于最小化负对数似然）。

伪代码

def LLaVA_Forward(image, input_ids):# 将图像通过视觉编码器得到视觉特征visual_features = VisualEncoder(image)# 通过投影层将视觉特征转换为语言模型相同的维度projected_features = ProjectionLayer(visual_features)# 将投影后的特征作为记忆传入语言模型（扩展成合适的维度）memory = ExpandDim(projected_features)# 语言模型接收文本输入和视觉记忆，输出各 token 的 logitslogits = LanguageModel(input_ids, memory)return logitsdef Compute_Loss(logits, target_ids, answer_mask):# 对 logits、target_ids 和 answer_mask 进行扁平化处理logits_flat = Flatten(logits)targets_flat = Flatten(target_ids)mask_flat = Flatten(answer_mask)  # 1 表示属于助手回答的 token# 仅对 mask 为1的 token 计算交叉熵损失selected_logits = logits_flat[mask_flat == 1]selected_targets = targets_flat[mask_flat == 1]loss = CrossEntropyLoss(selected_logits, selected_targets)return loss# 训练过程伪代码
for each batch in training_data:image, input_ids, target_ids, answer_mask = GetBatch()logits = LLaVA_Forward(image, input_ids)loss = Compute_Loss(logits, target_ids, answer_mask)loss.backward()optimizer.step()optimizer.zero_grad()

Instruction FineTuning案例代码

# 环境 torch transformers datasets Pillow sentencepiece accelerate
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForCausalLM, ViTModel, ViTImageProcessor
from torch.optim import AdamW
from datasets import load_dataset
from PIL import Image# --------------------
# 配置参数
# --------------------
class Config:dataset_name = "nlphuji/flickr30k"max_length = 128  # 总序列长度（图像+文本）image_size = 224text_model = "distilgpt2"vision_model = "google/vit-base-patch16-224"projector_hidden_size = 768batch_size = 8lr = 1e-5epochs = 3device = "cuda" if torch.cuda.is_available() else "cpu"# --------------------
# 数据集构建
# --------------------
class Flickr30kDataset(Dataset):def __init__(self, dataset, processor, tokenizer):self.dataset = datasetself.processor = processorself.tokenizer = tokenizerself.prompt_template = "[用户] 请描述这张图片。\n[助手] "self.text_max_length = Config.max_length - 1  # 为图像特征预留位置def __len__(self):return len(self.dataset)def __getitem__(self, idx):item = self.dataset[idx]image = item["image"].convert("RGB")text = item["caption"][0]# 图像处理pixel_values = self.processor(image, return_tensors="pt",size={"height": Config.image_size, "width": Config.image_size}).pixel_values[0]# 文本处理（预留图像位置）full_text = self.prompt_template + texttokens = self.tokenizer(full_text, max_length=self.text_max_length,truncation=True,padding="max_length",return_tensors="pt",return_attention_mask=True)# 构建助手标签掩码input_ids = tokens["input_ids"][0]prompt_ids = self.tokenizer.encode(self.prompt_template, add_special_tokens=False,max_length=self.text_max_length,truncation=True)prompt_len = len(prompt_ids)# 扩展标签到总长度（图像+文本）labels = torch.full((Config.max_length,), -100, dtype=torch.long)labels[1:1+len(input_ids)] = input_ids.clone()labels[1:1+prompt_len+1] = -100  # +1覆盖[助手]标记return {"pixel_values": pixel_values,"input_ids": input_ids,"labels": labels}# --------------------
# 多模态模型
# --------------------
class MultimodalModel(torch.nn.Module):def __init__(self):super().__init__()self.vision_encoder = ViTModel.from_pretrained(Config.vision_model)for param in self.vision_encoder.parameters():param.requires_grad = Falseself.projector = torch.nn.Linear(self.vision_encoder.config.hidden_size,Config.projector_hidden_size)self.language_model = AutoModelForCausalLM.from_pretrained(Config.text_model)self.language_model.transformer.wte.weight.requires_grad = Falsedef forward(self, pixel_values, input_ids, labels):# 视觉特征提取vision_outputs = self.vision_encoder(pixel_values)image_embeds = vision_outputs.last_hidden_state[:, 0, :]image_embeds = self.projector(image_embeds)# 文本嵌入text_embeds = self.language_model.transformer.wte(input_ids)# 拼接图像特征（维度对齐）inputs_embeds = torch.cat([image_embeds.unsqueeze(1),  # [batch, 1, hidden]text_embeds                 # [batch, seq-1, hidden]], dim=1)  # 最终维度 [batch, seq, hidden]# 语言模型前向outputs = self.language_model(inputs_embeds=inputs_embeds,labels=labels)return outputs# --------------------
# 训练流程
# --------------------
def train():processor = ViTImageProcessor.from_pretrained(Config.vision_model)tokenizer = AutoTokenizer.from_pretrained(Config.text_model)tokenizer.pad_token = tokenizer.eos_token# 加载数据集dataset = load_dataset(Config.dataset_name)["test"].shuffle(seed=42).select(range(100))train_dataset = Flickr30kDataset(dataset, processor, tokenizer)# 维度验证sample = train_dataset[0]assert sample["pixel_values"].shape == (3, Config.image_size, Config.image_size)assert sample["input_ids"].shape == (Config.max_length-1,)assert sample["labels"].shape == (Config.max_length,)train_loader = DataLoader(train_dataset,batch_size=Config.batch_size,shuffle=True,collate_fn=lambda batch: {'pixel_values': torch.stack([x['pixel_values'] for x in batch]),'input_ids': torch.stack([x['input_ids'] for x in batch]),'labels': torch.stack([x['labels'] for x in batch])})model = MultimodalModel().to(Config.device)optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=Config.lr)model.train()for epoch in range(Config.epochs):total_loss = 0for batch in train_loader:pixel_values = batch["pixel_values"].to(Config.device)input_ids = batch["input_ids"].to(Config.device)labels = batch["labels"].to(Config.device)outputs = model(pixel_values, input_ids, labels)loss = outputs.lossoptimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()print(f"Epoch: {epoch+1}, Step Loss: {loss.item():.4f}")print(f"Epoch {epoch+1} Average Loss: {total_loss/len(train_loader):.4f}")return model# --------------------
# 推理模块
# --------------------
def inference(model, image_path):processor = ViTImageProcessor.from_pretrained(Config.vision_model)tokenizer = AutoTokenizer.from_pretrained(Config.text_model)image = Image.open(image_path).convert("RGB")pixel_values = processor(image,return_tensors="pt",size={"height": Config.image_size, "width": Config.image_size}).pixel_values.to(Config.device)# 生成输入序列（适配维度）prompt = "[用户] 请描述这张图片。\n[助手] "input_ids = tokenizer.encode(prompt,max_length=Config.max_length-1,truncation=True,padding="max_length",return_tensors="pt").to(Config.device)model.eval()with torch.no_grad():vision_outputs = model.vision_encoder(pixel_values)image_embeds = model.projector(vision_outputs.last_hidden_state[:, 0, :])text_embeds = model.language_model.transformer.wte(input_ids)inputs_embeds = torch.cat([image_embeds.unsqueeze(1), text_embeds], dim=1)generated = model.language_model.generate(inputs_embeds=inputs_embeds,max_length=Config.max_length,temperature=0.9)print("生成结果:", tokenizer.decode(generated[0], skip_special_tokens=True))if __name__ == "__main__":trained_model = train()# torch.save(trained_model.state_dict(), "multimodal_model.pth")# inference(trained_model, "test_image.jpg")  # 替换为真实图片路径