CLIP 深度剖析：对比语言-图像预训练的完整技术栈

1
def info_nce_loss(image_embeds, text_embeds, temperature):
2
    """
3
    InfoNCE 对比损失（简化版）。
4

5
    参数:
6
        image_embeds: (B, D) 图像特征
7
        text_embeds:  (B, D) 文本特征
8
        temperature:  温度系数 τ
9
    """
10
    # 归一化
11
    image_embeds = F.normalize(image_embeds, dim=-1)
12
    text_embeds  = F.normalize(text_embeds, dim=-1)
13

14
    # 计算余弦相似度矩阵 (B, B)
15
    logits = torch.matmul(image_embeds, text_embeds.T) / temperature
16

17
    # labels: 对角线上是正样本
18
    batch_size = image_embeds.shape[0]
19
    labels = torch.arange(batch_size, device=image_embeds.device)
20

21
    # 两方向的对称损失
22
    loss_i2t = F.cross_entropy(logits, labels)       # 图像→文本
23
    loss_t2i = F.cross_entropy(logits.T, labels)      # 文本→图像
24

25
    loss = (loss_i2t + loss_t2i) / 2
26
    return loss

1
def temperature_effect():
2
    """
3
    温度系数 τ 对对比损失的影响。
4
    τ 是可学习的参数，初始化通常在 0.01~0.1 之间。
5
    """
6
    configs = {
7
        "τ 极小 (0.01)": "softmax 极度尖锐 → 只关注最难的负样本",
8
        "τ 适中 (0.07)": "平衡——正确配对突出，负样本被适度压制",
9
        "τ 过大 (1.0)":  "softmax 趋于均匀 → 无法区分正负样本",
10
    }
11
    return configs
12

13
    # CLIP 原始实现中 τ 被建模为 log(scaled_logits) 的可学习参数
14
    # logit_scale = nn.Parameter(torch.ones([]) * np.log(1/0.07))
15
    # temperature = 1 / torch.exp(logit_scale) ≈ 0.07

1
class ImagePatcher(nn.Module):
2
    """
3
    图像分块：把图像切成 patch 并线性投影。
4
    """
5
    def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768):
6
        super().__init__()
7
        self.image_size = image_size
8
        self.patch_size = patch_size
9
        self.num_patches = (image_size // patch_size) ** 2  # 196
10

11
        # 线性投影层（等价于 2D conv with kernel_size=patch_size）
12
        self.proj = nn.Conv2d(in_channels, embed_dim,
13
                               kernel_size=patch_size, stride=patch_size)
14

15
    def forward(self, x):
16
        """
17
        参数:
18
            x: (B, 3, 224, 224) 图像 batch
19
        返回:
20
            patches: (B, num_patches, embed_dim)  每个 patch 的嵌入
21
        """
22
        # Conv2d: (B, 3, 224, 224) → (B, embed_dim, 14, 14)
23
        x = self.proj(x)
24
        # 展平空间维度: (B, embed_dim, 14, 14) → (B, embed_dim, 196)
25
        x = x.flatten(2)
26
        # 调换维度: (B, embed_dim, 196) → (B, 196, embed_dim)
27
        x = x.transpose(1, 2)
28
        return x

1
class ViTEncoder(nn.Module):
2
    """
3
    CLIP 的 ViT 视觉编码器。
4
    """
5
    def __init__(self, image_size=224, patch_size=16, embed_dim=768,
6
                 num_layers=12, num_heads=12, mlp_ratio=4.0):
7
        super().__init__()
8
        self.patch_embed = ImagePatcher(image_size, patch_size, 3, embed_dim)
9
        self.num_patches = self.patch_embed.num_patches
10

11
        # [CLS] token
12
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
13

14
        # 位置编码（可学习）
15
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim))
16

17
        # Transformer 编码器
18
        encoder_layer = nn.TransformerEncoderLayer(
19
            d_model=embed_dim,
20
            nhead=num_heads,
21
            dim_feedforward=int(embed_dim * mlp_ratio),
22
            activation="gelu",
23
            batch_first=True,
24
            norm_first=True,  # Pre-LN (CLIP 使用)
25
        )
26
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
27

28
        # LayerNorm
29
        self.ln = nn.LayerNorm(embed_dim)
30

31
    def forward(self, x):
32
        """
33
        参数:
34
            x: (B, 3, 224, 224)
35
        返回:
36
            pooled: (B, embed_dim) [CLS] token 的最终表示
37
            sequence: (B, num_patches+1, embed_dim) 完整序列
38
        """
39
        B = x.shape[0]
40

41
        # 1) 分块 + 投影
42
        patches = self.patch_embed(x)  # (B, 196, embed_dim)
43

44
        # 2) 添加 [CLS] token
45
        cls_tokens = self.cls_token.expand(B, -1, -1)  # (B, 1, embed_dim)
46
        x = torch.cat([cls_tokens, patches], dim=1)   # (B, 197, embed_dim)
47

48
        # 3) 添加位置编码
49
        x = x + self.pos_embed  # (B, 197, embed_dim)
50

51
        # 4) Transformer 编码
52
        x = self.transformer(x)  # (B, 197, embed_dim)
53

54
        # 5) LayerNorm
55
        x = self.ln(x)
56

57
        # [CLS] token
58
        pooled = x[:, 0, :]       # (B, embed_dim)
59
        sequence = x              # (B, 197, embed_dim)
60

61
        return pooled, sequence

1
class TextTokenizer:
2
    """
3
    CLIP 使用 BPE (Byte Pair Encoding) 分词器。
4
    与 GPT-2 / RoBERTa 相同。
5
    """
6
    def __init__(self, vocab_size=50257):
7
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
8
        self.tokenizer.add_special_tokens(["<|startoftext|>", "<|endoftext|>"])
9

10
    def encode(self, texts):
11
        """
12
        参数:
13
            texts: str 或 list[str]
14
        返回:
15
            input_ids: (B, L) token IDs
16
            attention_mask: (B, L) 注意力掩码
17
        """
18
        if isinstance(texts, str):
19
            texts = [texts]
20

21
        encoded = self.tokenizer(
22
            texts,
23
            padding=True,
24
            truncation=True,
25
            max_length=77,     # CLIP 固定文本长度
26
            return_tensors="pt",
27
        )
28
        return encoded["input_ids"], encoded["attention_mask"]

1
class TextTransformerEncoder(nn.Module):
2
    """
3
    CLIP 的文本编码器。
4
    Transformer Encoder 处理文本 token 序列。
5
    """
6
    def __init__(self, vocab_size=50257, embed_dim=512,
7
                 context_length=77, num_layers=12, num_heads=8):
8
        super().__init__()
9
        self.context_length = context_length
10

11
        # Token 嵌入
12
        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
13

14
        # 位置编码（可学习，固定长度）
15
        self.pos_embedding = nn.Parameter(
16
            torch.zeros(1, context_length, embed_dim)  # (1, 77, 512)
17
        )
18

19
        # Transformer 编码器
20
        encoder_layer = nn.TransformerEncoderLayer(
21
            d_model=embed_dim,
22
            nhead=num_heads,
23
            dim_feedforward=embed_dim * 4,
24
            activation="gelu",
25
            batch_first=True,
26
            norm_first=True,
27
        )
28
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
29

30
        # LayerNorm
31
        self.ln = nn.LayerNorm(embed_dim)
32

33
    def forward(self, input_ids, attention_mask=None):
34
        """
35
        参数:
36
            input_ids:     (B, L) token IDs
37
            attention_mask: (B, L) 可选
38
        返回:
39
            pooled: (B, embed_dim) [EOS] token 的表示
40
            sequence: (B, L, embed_dim) 完整序列
41
        """
42
        B, L = input_ids.shape
43

44
        # 1) Token 嵌入
45
        x = self.token_embedding(input_ids)  # (B, L, embed_dim)
46

47
        # 2) 添加位置编码
48
        x = x + self.pos_embedding[:, :L, :]  # (B, L, embed_dim)
49

50
        # 3) Transformer 编码
51
        x = self.transformer(x)  # (B, L, embed_dim)
52

53
        # 4) LayerNorm
54
        x = self.ln(x)
55

56
        # 5) 取 [EOS] token（通常是文本序列最后一个 token）作为句子表示
57
        #    或者用 attention pooling
58
        pooled = x[:, -1, :]  # (B, embed_dim)
59

60
        return pooled, x

1
class WITDataset(Dataset):
2
    """
3
    CLIP 的训练数据加载器。
4
    """
5
    def __init__(self, image_paths, texts, image_size=224):
6
        self.image_paths = image_paths
7
        self.texts = texts
8
        self.image_size = image_size
9
        self.img_transform = transforms.Compose([
10
            transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
11
            transforms.CenterCrop(image_size),
12
            transforms.ToTensor(),
13
            transforms.Normalize(
14
                mean=[0.48145466, 0.4578275, 0.40821073],  # ImageNet RGB 统计
15
                std=[0.26862954, 0.26130258, 0.27577711],
16
            ),
17
        ])
18

19
    def __getitem__(self, idx):
20
        # 加载图像
21
        image = Image.open(self.image_paths[idx]).convert("RGB")
22
        image = self.img_transform(image)
23

24
        # 获取文本
25
        text = self.texts[idx][:77]  # 截断到 77 tokens
26

27
        return {"image": image, "text": text}

1
@dataclass
2
class CLIPTrainingConfig:
3
    """CLIP 训练配置。"""
4
    # 模型
5
    vision_model: str = "ViT-L/14"       # CLIP 默认使用 ViT-L/14
6
    embed_dim: int = 768                 # 视觉/文本特征的维度
7

8
    # 训练
9
    batch_size: int = 32768               # CLIP 原始实现用 32768（很大！）
10
    num_epochs: int = 32
11
    learning_rate: float = 5e-4            # AdamW 学习率
12
    weight_decay: float = 0.1
13
    warmup_steps: int = 5000
14

15
    # 优化器
16
    optimizer: str = "AdamW"
17
    beta1: float = 0.9
18
    beta2: float = 0.98
19
    eps: float = 1e-6
20

21
    # 损失
22
    temperature_init: float = 0.07        # 温度系数初始化
23
    label_smoothing: float = 0.0
24

25
    # 正则化
26
    gradient_clip: float = 1.0
27
    mixed_precision: bool = True         # FP16 混合精度
28

29
    # 数据
30
    image_size: int = 224
31
    context_length: int = 77
32
    max_text_length: int = 77

1
def train_clip(model, dataloader, config):
2
    """
3
    CLIP 训练循环（简化版）。
4
    """
5
    optimizer = torch.optim.AdamW(
6
        model.parameters(),
7
        lr=config.learning_rate,
8
        weight_decay=config.weight_decay,
9
        betas=(config.beta1, config.beta2),
10
        eps=config.eps,
11
    )
12

13
    # 学习率调度器
14
    def lr_lambda(step):
15
        if step < config.warmup_steps:
16
            return step / config.warmup_steps
17
        else:
18
            progress = (step - config.warmup_steps) / (config.total_steps - config.warmup_steps)
19
            return 0.5 * (1 + np.cos(np.pi * progress))
20

21
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
22

23
    # 可学习的温度系数
24
    logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / config.temperature_init))
25
    logit_scale = logit_scale.to(device)
26

27
    model.train()
28
    for epoch in range(config.num_epochs):
29
        epoch_loss = 0.0
30
        for batch in tqdm(dataloader):
31
            images = batch["image"].to(device)
32
            texts  = batch["text"]
33

34
            # 编码
35
            image_features = model.visual(images)    # (B, D)
36
            text_features  = model.text(texts)       # (B, D)
37

38
            # L2 归一化
39
            image_features = F.normalize(image_features, dim=-1)
40
            text_features  = F.normalize(text_features, dim=-1)
41

42
            # 对比 logits
43
            temperature = (1 / torch.exp(logit_scale))
44
            logits_per_image = (image_features @ text_features.T) / temperature  # (B, B)
45
            logits_per_text  = logits_per_image.T                                        # (B, B)
46

47
            # Labels: 对角线是正样本
48
            batch_size = images.shape[0]
49
            labels = torch.arange(batch_size, device=device)
50

51
            # 对称交叉熵损失
52
            loss_i2t = F.cross_entropy(logits_per_image, labels)
53
            loss_t2i = F.cross_entropy(logits_per_text, labels)
54
            loss = (loss_i2t + loss_t2i) / 2
55

56
            # 反向传播
57
            optimizer.zero_grad()
58
            loss.backward()
59

60
            # 梯度裁剪
61
            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
62

63
            optimizer.step()
64
            scheduler.step()
65

66
            epoch_loss += loss.item()
67

68
        # 打印
69
        print(f"Epoch {epoch}: loss={epoch_loss / len(dataloader):.4f}, "
70
              f"temp={temperature.item():.4f}")

1
@torch.no_grad()
2
def zero_shot_classify(model, image, class_names, device):
3
    """
4
    CLIP 零样本分类。
5

6
    参数:
7
        model:       CLIP 模型
8
        image:       (C, H, W) 图像 tensor
9
        class_names: list[str] 类别名称，如 ["cat", "dog", "bird"]
10
        device:      计算设备
11
    返回:
12
        预测的类别索引和概率
13
    """
14
    model.eval()
15

16
    # 1) 构建文本提示（class name → sentence）
17
    #    CLIP 原论文发现"a photo of a {class}"效果最好
18
    prompts = [f"a photo of a {c}." for c in class_names]
19

20
    # 2) 编码文本
21
    text_tokens = model.tokenizer(prompts).to(device)
22
    text_features = model.encode_text(text_tokens)
23
    text_features = F.normalize(text_features, dim=-1)  # (num_classes, D)
24

25
    # 3) 编码图像
26
    if len(image.shape) == 3:
27
        image = image.unsqueeze(0)  # (1, C, H, W)
28
    image = image.to(device)
29
    image_features = model.encode_image(image)
30
    image_features = F.normalize(image_features, dim=-1)  # (1, D)
31

32
    # 4) 计算余弦相似度
33
    similarity = (image_features @ text_features.T)  # (1, num_classes)
34
    probs = F.softmax(similarity * model.logit_scale, dim=-1)  # (1, num_classes)
35

36
    # 5) 返回结果
37
    probs = probs[0].cpu().numpy()
38
    pred_idx = probs.argmax()
39
    return pred_idx, class_names[pred_idx], probs[pred_idx], probs

1
def prompt_ensemble(model, class_names, image, device):
2
    """
3
    提示集成：使用多个提示模板提升零样本分类精度。
4
    CLIP 原论文发现这一招效果显著。
5
    """
6
    # 多种提示模板
7
    prompt_templates = [
8
        "a photo of a {}.",
9
        "a blurry photo of a {}.",
10
        "a low resolution photo of a {}.",
11
        "a bright photo of a {}.",
12
        "a dark photo of a {}.",
13
        "a photo of a small {}.",
14
        "a photo of a large {}.",
15
        "a photo of the {}.",
16
        "a photo of a {} in the scene.",
17
        "a photo of the {} in the scene.",
18
    ]
19

20
    all_text_features = []
21

22
    for template in prompt_templates:
23
        prompts = [template.format(c) for c in class_names]
24
        text_tokens = model.tokenizer(prompts).to(device)
25
        text_features = model.encode_text(text_tokens)
26
        text_features = F.normalize(text_features, dim=-1)
27
        all_text_features.append(text_features)
28

29
    # 平均所有提示的特征
30
    text_features_avg = torch.stack(all_text_features, dim=0).mean(dim=0)  # (num_classes, D)
31
    text_features_avg = F.normalize(text_features_avg, dim=-1)
32

33
    # 图像编码
34
    image_features = model.encode_image(image.to(device).unsqueeze(0))
35
    image_features = F.normalize(image_features, dim=-1)
36

37
    # 相似度
38
    similarity = image_features @ text_features_avg.T
39
    probs = F.softmax(similarity * model.logit_scale, dim=-1)
40

41
    return probs[0].cpu().numpy()

1
def siglip_loss(image_features, text_features, temperature):
2
    """
3
    SigLIP 的 Sigmoid 对比损失。
4

5
    特点:
6
      - 每个样本独立优化（不需要 global softmax）
7
      - 可以处理更大的 batch size
8
      - 避免正样本被其他正样本"稀释"
9
    """
10
    image_features = F.normalize(image_features, dim=-1)
11
    text_features  = F.normalize(text_features, dim=-1)
12

13
    logits = image_features @ text_features.T / temperature  # (B, B)
14

15
    # 正样本标签
16
    B = image_features.shape[0]
17
    labels = torch.eye(B, device=image_features.device)
18

19
    # Sigmoid 交叉熵（而非 softmax）
20
    loss = F.binary_cross_entropy_with_logits(logits, labels)
21

22
    return loss