07-多模态学习¶

import torch
import torch.nn as nn

class CrossAttention(nn.Module):  # 继承nn.Module定义神经网络层
    """跨模态注意力：query来自模态A, key/value来自模态B"""
    def __init__(self, dim, n_heads=8):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        self.n_heads = n_heads
        self.head_dim = dim // n_heads

        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)
        self.scale = self.head_dim ** -0.5

    def forward(self, query, context):
        """
        query: [B, Lq, D] 模态A的序列
        context: [B, Lk, D] 模态B的序列
        """
        B, Lq, D = query.shape
        Lk = context.size(1)

        Q = self.q_proj(query).view(B, Lq, self.n_heads, self.head_dim).transpose(1, 2)  # 链式调用，连续执行多个方法
        K = self.k_proj(context).view(B, Lk, self.n_heads, self.head_dim).transpose(1, 2)
        V = self.v_proj(context).view(B, Lk, self.n_heads, self.head_dim).transpose(1, 2)

        attn = (Q @ K.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)

        out = (attn @ V).transpose(1, 2).reshape(B, Lq, D)  # reshape重塑张量形状
        return self.out_proj(out)

class MultimodalFusionBlock(nn.Module):
    """双向跨模态融合块"""
    def __init__(self, dim):
        super().__init__()
        self.text_to_image = CrossAttention(dim)
        self.image_to_text = CrossAttention(dim)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn_image = nn.Sequential(
            nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)
        )
        self.ffn_text = nn.Sequential(
            nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)
        )

    def forward(self, image_feat, text_feat):
        # 图像关注文本
        image_feat = image_feat + self.text_to_image(image_feat, text_feat)
        image_feat = image_feat + self.ffn_image(self.norm1(image_feat))

        # 文本关注图像
        text_feat = text_feat + self.image_to_text(text_feat, image_feat)
        text_feat = text_feat + self.ffn_text(self.norm2(text_feat))

        return image_feat, text_feat

class SimpleCaptioner(nn.Module):
    """简化的图像描述生成模型"""
    def __init__(self, visual_dim, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.visual_proj = nn.Linear(visual_dim, embed_dim)
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(embed_dim, nhead=8),
            num_layers=6
        )
        self.output_proj = nn.Linear(embed_dim, vocab_size)

    def forward(self, visual_features, token_ids):
        """
        visual_features: [B, N_patch, visual_dim] 图像特征
        token_ids: [B, L] 文本token序列
        """
        # 投影视觉特征
        memory = self.visual_proj(visual_features)  # [B, N_patch, embed_dim]
        memory = memory.transpose(0, 1)  # [N_patch, B, embed_dim]

        # 文本嵌入
        tgt = self.embedding(token_ids).transpose(0, 1)  # [L, B, embed_dim]

        # 因果掩码
        L = tgt.size(0)
        causal_mask = torch.triu(torch.ones(L, L), diagonal=1).bool()

        # 解码
        out = self.decoder(tgt, memory, tgt_mask=causal_mask)
        return self.output_proj(out.transpose(0, 1))  # [B, L, vocab_size]

import torch
import torch.nn.functional as F
from PIL import Image

# 使用HuggingFace的transformers加载CLIP
# pip install transformers
from transformers import CLIPProcessor, CLIPModel

def clip_zero_shot_classification(image_path, candidate_labels):
    """CLIP零样本图像分类"""
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    image = Image.open(image_path)
    texts = [f"a photo of a {label}" for label in candidate_labels]  # 列表推导式，简洁创建列表

    inputs = processor(text=texts, images=image, return_tensors="pt", padding=True)

    with torch.no_grad():  # 禁用梯度计算，节省内存
        outputs = model(**inputs)
        logits = outputs.logits_per_image  # [1, n_labels]
        probs = logits.softmax(dim=-1)

    for label, prob in zip(candidate_labels, probs[0]):  # zip按位置配对多个可迭代对象
        print(f"  {label}: {prob.item():.4f}")  # .item()将单元素张量转为Python数值

    predicted = candidate_labels[probs.argmax().item()]
    return predicted

def clip_image_text_retrieval(image_paths, query_text):
    """CLIP文搜图检索"""
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    images = [Image.open(p) for p in image_paths]
    inputs = processor(text=[query_text], images=images,
                      return_tensors="pt", padding=True)

    with torch.no_grad():
        outputs = model(**inputs)
        similarities = outputs.logits_per_text[0]  # [n_images]

    ranked = similarities.argsort(descending=True)
    print(f"查询: '{query_text}'")
    for rank, idx in enumerate(ranked):  # enumerate同时获取索引和元素
        print(f"  #{rank+1}: {image_paths[idx]} (score: {similarities[idx]:.3f})")

    return ranked

# 使用示例
# clip_zero_shot_classification("cat.jpg", ["cat", "dog", "bird", "car"])
# clip_image_text_retrieval(["a.jpg", "b.jpg", "c.jpg"], "a dog playing frisbee")