多模态学习（三）：多模态大模型与统一架构

# 使用 COCO API 加载和可视化图像描述
from pycocotools.coco import COCO
import matplotlib.pyplot as plt

coco = COCO("annotations/captions_val2017.json")
img_ids = coco.getImgIds()[:3]
for img_id in img_ids:
    img_info = coco.loadImgs(img_id)[0]
    ann_ids = coco.getAnnIds(imgIds=img_id)
    anns = coco.loadAnns(ann_ids)
    captions = [ann["caption"] for ann in anns]
    print(f"Image {img_info['file_name']}:")
    for i, cap in enumerate(captions[:3], 1):
        print(f"  [{i}] {cap}")

# 构建图像-文本对的数据加载器
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as T

class CaptionDataset(Dataset):
    def __init__(self, image_dir, annotations, vocab):
        self.image_dir = image_dir
        self.annotations = annotations
        self.vocab = vocab
        self.transform = T.Compose([
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406],
                        [0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, idx):
        ann = self.annotations[idx]
        image = Image.open(f"{self.image_dir}/{ann['image']}")
        image = self.transform(image)
        caption_ids = self.vocab.encode(ann["caption"])
        return image, caption_ids

# Show and Tell 风格的 Encoder-Decoder 模型
import torch
import torch.nn as nn
import torchvision.models as models

class ShowAndTell(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        # Encoder: 预训练 ResNet-50
        resnet = models.resnet50(weights="IMAGENET1K_V1")
        self.encoder = nn.Sequential(*list(resnet.children())[:-1])
        for param in self.encoder.parameters():
            param.requires_grad = False

        # Decoder: LSTM
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, images, captions):
        features = self.encoder(images).squeeze(-1).squeeze(-1)
        embeddings = self.embed(captions)
        h0, c0 = self._init_hidden(features)
        outputs, _ = self.lstm(embeddings, (h0, c0))
        return self.fc(outputs)

    def _init_hidden(self, features):
        h0 = features.unsqueeze(0)
        c0 = torch.zeros_like(h0)
        return h0, c0

# 基于 ViT 的视觉 Encoder
import torch.nn as nn
from transformers import ViTModel

class ViTEncoder(nn.Module):
    def __init__(self, model_name="google/vit-base-patch16-224"):
        super().__init__()
        self.vit = ViTModel.from_pretrained(model_name)
        for param in self.vit.parameters():
            param.requires_grad = False

    def forward(self, images):
        outputs = self.vit(pixel_values=images)
        # 使用 [CLS] token 作为全局表示
        cls_token = outputs.last_hidden_state[:, 0, :]
        # 使用所有 patch token 作为空间特征
        patch_tokens = outputs.last_hidden_state[:, 1:, :]
        return cls_token, patch_tokens

# 测试
encoder = ViTEncoder()
dummy = torch.randn(2, 3, 224, 224)
cls_feat, patch_feat = encoder(dummy)
print(f"CLS: {cls_feat.shape}")       # [2, 768]
print(f"Patches: {patch_feat.shape}")  # [2, 196, 768]

# 软注意力（Soft Attention）机制实现
class SoftAttention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim):
        super().__init__()
        self.encoder_att = nn.Linear(encoder_dim, decoder_dim)
        self.decoder_att = nn.Linear(decoder_dim, decoder_dim)
        self.full_att = nn.Linear(decoder_dim, 1)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, encoder_out, decoder_hidden):
        # encoder_out: (batch, num_pixels, encoder_dim)
        att1 = self.encoder_att(encoder_out)        # (batch, num_pixels, decoder_dim)
        att2 = self.decoder_att(decoder_hidden)      # (batch, decoder_dim)
        att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2)
        alpha = self.softmax(att)  # (batch, num_pixels)
        attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)
        return attention_weighted_encoding, alpha

# 自顶向下注意力（Top-Down Attention）
class TopDownAttention(nn.Module):
    def __init__(self, obj_dim, hidden_dim, num_heads=8):
        super().__init__()
        self.multihead_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            batch_first=True,
            kdim=obj_dim,
            vdim=obj_dim,
        )
        self.obj_proj = nn.Linear(obj_dim, hidden_dim)

    def forward(self, decoder_query, object_features, mask=None):
        key = self.obj_proj(object_features)
        value = object_features
        attended, attn_weights = self.multihead_attn(
            query=decoder_query,
            key=key,
            value=value,
            key_padding_mask=mask,
        )
        return attended, attn_weights

# object_features 来自 Faster R-CNN 的 RoI Align
# shape: (batch, num_objects, obj_dim)

# 使用 Stable Diffusion 生成图像
from diffusers import StableDiffusionPipeline
import torch

# 加载预训练模型
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

# 文本到图像生成
prompt = "a photograph of an astronaut riding a horse"
image = pipe(
    prompt,
    num_inference_steps=50,
    guidance_scale=7.5,
).images[0]
image.save("astronaut_horse.png")

# DDPM 扩散过程的核心公式实现
import torch

class DiffusionProcess:
    def __init__(self, num_steps=1000):
        self.num_steps = num_steps
        beta = torch.linspace(1e-4, 0.02, num_steps)
        self.alpha = 1.0 - beta
        self.alpha_bar = torch.cumprod(self.alpha, dim=0)

    def add_noise(self, x0, t):
        noise = torch.randn_like(x0)
        sqrt_alpha_bar = torch.sqrt(self.alpha_bar[t])
        sqrt_one_minus_alpha_bar = torch.sqrt(1 - self.alpha_bar[t])
        return sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * noise, noise

    def sample(self, model, xT, text_embed):
        x = xT
        for t in reversed(range(self.num_steps)):
            noise_pred = model(x, t, text_embed)
            # 简化的反向采样步骤
            x = self._reverse_step(x, noise_pred, t)
        return x

# 计算 BLEU 和 CIDEr 分数
from pycocoevalcap.bleu.bleu import Bleu
from pycocoevalcap.cider.cider import Cider
from pycocoevalcap.tokenizer.ptbtokenizer import PTBTokenizer

# 准备数据格式
gts = {"img1": [{"caption": "a dog playing in the snow"}]}
res = {"img1": [{"caption": "a dog is playing in white snow"}]}

# Tokenize
tokenizer = PTBTokenizer()
gts = tokenizer.tokenize(gts)
res = tokenizer.tokenize(res)

# BLEU
bleu_scorer = Bleu(4)
bleu, _ = bleu_scorer.compute_score(gts, res)
for i, score in enumerate(bleu, 1):
    print(f"BLEU-{i}: {score:.4f}")

# CIDEr
cider_scorer = Cider()
cider, _ = cider_scorer.compute_score(gts, res)
print(f"CIDEr: {cider:.4f}")

# 使用 SPICE 进行语义级评估
from pycocoevalcap.spice.spice import Spice

# SPICE 需要 Java 环境
spice_scorer = Spice()
score, _ = spice_scorer.compute_score(gts, res)
print(f"SPICE: {score:.4f}")

# CLIPScore: 基于 CLIP 的无参考评估
import clip
from PIL import Image
import torch.nn.functional as F

model, preprocess = clip.load("ViT-B/32")
image = preprocess(Image.open("photo.jpg")).unsqueeze(0)
text = clip.tokenize(["a dog playing in the snow"])

with torch.no_grad():
    img_feat = model.encode_image(image)
    txt_feat = model.encode_text(text)
    score = F.cosine_similarity(img_feat, txt_feat)
    print(f"CLIPScore: {score.item():.4f}")

# 使用 ControlNet 进行姿态控制生成
from diffusers import (
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    UniPCMultistepScheduler,
)
from controlnet_aux import OpenposeDetector
from PIL import Image
import torch

# 加载 ControlNet
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/control_v11p_sd15_openpose",
    torch_dtype=torch.float16,
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

# 提取姿态并生成图像
pose_detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
input_image = Image.open("person.jpg")
pose = pose_detector(input_image)
image = pipe(
    "a beautiful portrait, 4k quality",
    image=pose,
    num_inference_steps=20,
).images[0]

# 提示词工程：构建高质量提示
class PromptBuilder:
    def __init__(self):
        self.quality_terms = [
            "masterpiece", "best quality", "highres", "detailed"
        ]
        self.style_terms = {
            "photo": "photorealistic, natural lighting, 85mm lens",
            "anime": "anime style, cel shading, vibrant colors",
            "painting": "oil painting, impressionist style, canvas texture",
        }
        self.negative_default = (
            "worst quality, low quality, blurry, deformed, "
            "bad anatomy, extra limbs, watermark"
        )

    def build(self, subject, style="photo"):
        parts = self.quality_terms.copy()
        parts.append(self.style_terms.get(style, ""))
        parts.append(subject)
        return ", ".join(parts), self.negative_default

builder = PromptBuilder()
pos, neg = builder.build("a cat sitting on a windowsill", "photo")
print(f"Positive: {pos}")
print(f"Negative: {neg}")

# 使用 BLIP-2 进行零样本图像描述
from transformers import Blip2Processor, Blip2ForConditionalGeneration
from PIL import Image
import torch

processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
model = Blip2ForConditionalGeneration.from_pretrained(
    "Salesforce/blip2-opt-2.7b",
    torch_dtype=torch.float16,
    device_map="auto",
)

image = Image.open("scene.jpg")
inputs = processor(images=image, return_tensors="pt").to("cuda", torch.float16)
generated_ids = model.generate(**inputs, max_new_tokens=50)
caption = processor.batch_decode(generated_ids, skip_special_tokens=True)
print(f"Generated: {caption[0].strip()}")

# 带提示的问题生成
question = "What color is the sky?"
inputs_qa = processor(
    images=image, text=question, return_tensors="pt"
).to("cuda", torch.float16)
out = model.generate(**inputs_qa, max_new_tokens=20)
answer = processor.batch_decode(out, skip_special_tokens=True)
print(f"Q: {question}  A: {answer[0].strip()}")

# 使用 LoRA 微调 Stable Diffusion
from diffusers import StableDiffusionPipeline, UNet2DConditionModel
from peft import LoraConfig, get_peft_model

# 加载基础模型
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5"
)
unet = pipe.unet

# 配置 LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout=0.1,
    bias="none",
)
lora_unet = get_peft_model(unet, lora_config)
lora_unet.print_trainable_parameters()
# trainable params: 8,838,144 || all params: 859,520,964
# trainable%: 1.03%

# 训练循环中使用梯度检查点节省显存
unet.enable_gradient_checkpointing()
pipe.enable_xformers_memory_efficient_attention()

# 保存 LoRA 权重
lora_unet.save_pretrained("lora-sd-v1", safe_serialization=True)