CLIP模型实战入门：从零构建图像-文本匹配系统

1次阅读

共计 2683 个字符，预计需要花费 7 分钟才能阅读完成。

传统图像分类模型（如 ResNet）依赖固定类别的监督训练，而 CLIP（Contrastive Language-Image Pretraining）通过对比学习实现了图像和文本的联合嵌入空间对齐。这种范式突破了两大限制：

类别开放性：无需预定义标签集，直接处理自然语言描述
跨模态检索：支持以图搜文、以文搜图的双向查询

CLIP 的训练目标是最小化以下对比损失函数：

$$
\mathcal{L} = -\frac{1}{N}\left(\sum_{i=1}^N \log \frac{e^{\text{sim}(I_i,T_i)/\tau}}{\sum_{j=1}^N e^{\text{sim}(I_i,T_j)/\tau}} + \sum_{i=1}^N \log \frac{e^{\text{sim}(T_i,I_i)/\tau}}{\sum_{j=1}^N e^{\text{sim}(T_i,I_j)/\tau}}\right)
$$

其中：
– $\text{sim}(I,T)=I^TT/|I||T|$ 表示余弦相似度
– $\tau$ 为温度系数
– $N$ 是 batch 内样本数

正样本对（匹配的图文对）会被拉近，而负样本对（非匹配组合）则被推远。

import torch
import clip
from PIL import Image
from typing import Tuple, Optional

assert torch.__version__ >= '1.12.0', "需要 PyTorch 1.12+ 版本"

def load_model(device: str = "cuda" if torch.cuda.is_available() else "cpu") -> Tuple[torch.nn.Module, torch.nn.Module]:
    """加载预训练模型和预处理函数"""
    try:
        model, preprocess = clip.load("ViT-B/32", device=device)
        return model, preprocess
    except Exception as e:
        print(f"模型加载失败: {e}")
        raise

# 示例使用
model, preprocess = load_model()

def compute_similarity(
    image_path: str, 
    text_list: list[str],
    model: torch.nn.Module,
    preprocess,
    device: str
) -> torch.Tensor:
    """计算图像与多个文本的相似度"""
    try:
        # 图像编码
        image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
        image_features = model.encode_image(image)

        # 文本编码（自动处理截断）text_inputs = clip.tokenize(text_list, truncate=True).to(device)
        text_features = model.encode_text(text_inputs)

        # 归一化后计算余弦相似度
        image_features /= image_features.norm(dim=-1, keepdim=True)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        return (image_features @ text_features.T).squeeze(0)
    except FileNotFoundError:
        print("图像文件不存在")
        raise
    except RuntimeError as e:
        print(f"GPU 内存不足: {e}")
        raise

# 不同 batch_size 的显存占用测试
for bs in [32, 64, 128]:
    try:
        dummy_images = torch.randn(bs, 3, 224, 224).to(device)
        model.encode_image(dummy_images)
        print(f"Batch size {bs} 运行成功")
    except:
        print(f"Batch size {bs} 超出显存限制")

# 在模型定义中启用
from torch.utils.checkpoint import checkpoint

class CustomCLIP(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.model = original_model

    def forward(self, x):
        return checkpoint(self.model.encode_image, x)

CLIP 的 tokenizer 默认限制 77 个 token
中文等非空格分隔语言需要特别处理
解决方案：提前用 clip.tokenize(text, truncate=True) 自动截断

官方预处理会将图像强制缩放到 224×224
对于细粒度分类任务（如鸟类识别）可能丢失细节
改进方案：
使用 InterpolationMode.BICUBIC 重采样
保留原始比例做中心裁剪

# 安全启用混合精度
scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    image_features = model.encode_image(images)
    text_features = model.encode_text(texts)
    loss = contrastive_loss(image_features, text_features)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()