CLIP 多模态对比学习

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCLIP(nn.Module):
    def __init__(self, embed_dim=512):
        super().__init__()
        # 图像编码器（简化版：用 MLP 模拟）
        self.image_encoder = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Linear(512, embed_dim),
        )
        # 文本编码器（简化版：用 MLP 模拟）
        self.text_encoder = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Linear(512, embed_dim),
        )
        # 可学习的温度参数（log 缩放）
        self.log_temperature = nn.Parameter(torch.log(torch.tensor(0.07)))

    def forward(self, images, texts):
        # 编码
        img_features = self.image_encoder(images)    # (N, embed_dim)
        txt_features = self.text_encoder(texts)      # (N, embed_dim)

        # L2 归一化
        img_features = F.normalize(img_features, dim=-1)
        txt_features = F.normalize(txt_features, dim=-1)

        # 计算余弦相似度矩阵 (N × N)
        temperature = torch.exp(self.log_temperature)
        logits = torch.matmul(img_features, txt_features.T) / temperature

        # InfoNCE 损失（图像→文本 和 文本→图像 的对称损失）
        labels = torch.arange(logits.size(0), device=logits.device)
        loss_i2t = F.cross_entropy(logits, labels)     # 图像找文本
        loss_t2i = F.cross_entropy(logits.T, labels)   # 文本找图像
        loss = (loss_i2t + loss_t2i) / 2
        return loss

# 训练循环
model = SimpleCLIP()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(100):
    images = torch.randn(32, 768)  # 32 张图的特征
    texts  = torch.randn(32, 768)  # 32 段文本的特征
    loss = model(images, texts)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()