2. 视觉问答模型微调基础

任务说明

本案例将使用 Hugging Face Transformers 库,在 VQAv2 数据集子集上微调 Vision Encoder-Decoder 模型(ViT + ),使其能够根据输入图像回答自然语言问题。完成微调后,模型应能对未见过的图像-问题对生成合理答案。

环境依赖

pip install transformers datasets torch pillow accelerate

完整实现代码

import torch
from transformers import (
    VisionEncoderDecoderModel,
    ViTImageProcessor,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset
from PIL import Image
import numpy as np

# ========== 1. 加载模型与处理器 ==========
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit--image-captioning")
image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit--image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit--image-captioning")

# 设置特殊用于VQA任务
tokenizer.add_special_tokens({"pad_": "[PAD]"})
model.decoder.resize__embeddings(len(tokenizer))
model.config.decoder_start__id = tokenizer.bos__id
model.config.pad__id = tokenizer.pad__id
model.config.eos__id = tokenizer.eos__id

# ========== 2. 加载并预处理VQAv2子集 ==========
# 使用少量样本演示(完整训练请替换为完整数据集)
dataset = load_dataset("HuggingFaceM4/VQAv2", split="validation[:200]")

def preprocess_batch(batch):
    """将图像和问答对转换为模型输入格式"""
    inputs = []
    labels = []

    for img, question, answer in zip(batch["image"], batch["question"], batch["multiple_choice_answer"]):
        # 处理图像
        if img.mode != "RGB":
            img = img.convert("RGB")
        pixel_values = image_processor(img, return_tensors="pt").pixel_values

        # 构造输入文本:问题前加"Q:",答案前加"A:"
        question_text = f"Q: {question} A:"
        answer_text = f" {answer}"

        # 编码文本
        question_enc = tokenizer(question_text, return_tensors="pt", padding="max_length", max_length=32, truncation=True)
        answer_enc = tokenizer(answer_text, return_tensors="pt", padding="max_length", max_length=16, truncation=True)

        inputs.append({
            "pixel_values": pixel_values.squeeze(0),
            "labels": answer_enc.input_ids.squeeze(0),
            "decoder_input_ids": question_enc.input_ids.squeeze(0),
        })

    return {
        "pixel_values": torch.stack([x["pixel_values"] for x in inputs]),
        "labels": torch.stack([x["labels"] for x in inputs]),
        "decoder_input_ids": torch.stack([x["decoder_input_ids"] for x in inputs]),
    }

# 应用预处理
dataset = dataset.map(preprocess_batch, batched=True, remove_columns=dataset.column_names)
dataset.set_format(type="torch", columns=["pixel_values", "labels", "decoder_input_ids"])

# ========== 3. 配置训练参数 ==========
training_args = TrainingArguments(
    output_dir="./vqa_finetuned",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=2,
    learning_rate=5e-5,
    warmup_steps=100,
    logging_steps=10,
    save_steps=50,
    evaluation_strategy="no",
    save_total_limit=2,
    remove_unused_columns=False,
    fp16=torch.cuda.is_available(),
)

# ========== 4. 自定义Trainer以支持多模态输入 ==========
class MultimodalTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        pixel_values = inputs.pop("pixel_values")
        labels = inputs.pop("labels")
        decoder_input_ids = inputs.pop("decoder_input_ids")

        outputs = model(
            pixel_values=pixel_values,
            labels=labels,
            decoder_input_ids=decoder_input_ids,
        )
        loss = outputs.loss
        return (loss, outputs) if return_outputs else loss

# ========== 5. 训练 ==========
trainer = MultimodalTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

trainer.train()

# ========== 6. 保存模型 ==========
model.save_pretrained("./vqa_finetuned")
tokenizer.save_pretrained("./vqa_finetuned")
image_processor.save_pretrained("./vqa_finetuned")

运行方式

# 直接运行训练脚本
python train_vqa.py

# 训练完成后,使用以下推理代码测试模型
python inference_vqa.py --image_path test.jpg --question "What color is the car?"

推理验证代码

from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from PIL import Image
import torch

# 加载微调后的模型
model = VisionEncoderDecoderModel.from_pretrained("./vqa_finetuned")
image_processor = ViTImageProcessor.from_pretrained("./vqa_finetuned")
tokenizer = AutoTokenizer.from_pretrained("./vqa_finetuned")

def vqa_inference(image_path, question):
    # 加载并处理图像
    image = Image.open(image_path).convert("RGB")
    pixel_values = image_processor(image, return_tensors="pt").pixel_values

    # 构造输入文本
    input_text = f"Q: {question} A:"
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids

    # 生成答案
    with torch.no_grad():
        output_ids = model.generate(
            pixel_values=pixel_values,
            decoder_input_ids=input_ids,
            max_length=20,
            num_beams=4,
            early_stopping=True,
        )

    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return answer

# 测试示例
if __name__ == "__main__":
    result = vqa_inference("test_image.jpg", "What is in the image?")
    print(f"模型回答: {result}")

验证方法与通过标准

  1. 训练收敛验证:训练日志中loss应随step增加呈下降趋势,最终loss应低于初始loss的50%。
  2. 推理效果验证:使用训练集外的图像-问题对进行测试,模型应生成与图像内容相关的合理答案(如对猫的图像问"What animal is this?",应回答"cat"或相近词)。
  3. 定量评估(可选):在100条验证集样本上计算准确率(答案完全匹配),微调后准确率应高于未微调基线模型至少5个百分点。

常见问题

  • 显存不足:降低per_device_train_batch_size至2或1,或启用梯度累积。
  • 图像格式错误:确保所有输入图像为RGB三通道,若为RGBA或L模式需转换。
  • 答案格式不一致:VQAv2中答案可能包含多个词,训练时使用multiple_choice_answer字段保证答案规范性。