""" ============================================================================= CS525 · 第 2 讲共享模块：一个可以在笔记本上几秒钟训练完的"迷你 CLIP"（clip_lite.py） ============================================================================= 对应讲义：lectures/lecture02_reliable_generalization.md §4–5 ----------------------------------------------------------------------------- 【CLIP 在做什么】 ----------------------------------------------------------------------------- CLIP（Radford et al. 2021）有两个编码器：一个把图片变成向量，一个把文字变成向量。 训练时拿一批（图, 文）对，要求"配对的图和文"的向量相似、"不配对的"不相似——这叫对比学习（contrastive learning）， 损失函数叫 InfoNCE： 一个 batch 里有 B 对 (图_i, 文_i)。算出 B × B 的相似度矩阵 S[i, j] = τ · （向量已归一化）。 对每一行做 softmax，要求第 i 行的最大值在第 i 列（"这张图的说明是第 i 段文字"）；对每一列也这样做。 损失 = 两个方向的交叉熵的平均。 训练完之后，CLIP 可以**零样本（zero-shot）分类**：把每个类别名写成一句提示（"a photo of a husky"）， 用文字编码器得到 40 个类别向量；一张图属于哪个类，就看它的图片向量和哪个类别向量最像。 不需要任何带标签的训练图片。 ----------------------------------------------------------------------------- 【这里的迷你版】 ----------------------------------------------------------------------------- 图片编码器：线性层（48 → 32 维），或者一个隐层的 MLP（48 → 128 → 32） 文字编码器：词袋模型——每个词一个可学习的向量，取平均，再过一个线性层 温度 τ：可学习，初始 1/0.07（和 CLIP 一样），上限 100 此外还有两个"监督学习"的工具，用来做控制变量实验（第 2 讲代码 03）和微调（代码 02）： train_supervised 同样的图片编码器 + 一个线性分类头，用交叉熵训练（"标准的 ImageNet 训练"） fine_tune 从零样本模型出发，用带标签的数据接着训练（端到端微调） interpolate 两个模型的参数做线性插值（WiSE-FT，Wortsman et al. 2022） 【确定性】导入时调用 torch.set_num_threads(1)；每次训练用自己的 torch.Generator 抽 batch， 同一台机器上重复运行，结果逐字相同。 【运行自检】 python3 clip_lite.py # 约 10 秒：在网络图文对上训练一个迷你 CLIP，打印零样本准确率 """ import copy import numpy as np import torch import torch.nn as nn import torch.nn.functional as F torch.set_num_threads(1) # ============================================================================= # 模型 # ============================================================================= class ImageEncoder(nn.Module): """ 图片编码器：特征向量 → 嵌入向量（未归一化）。 【参数】in_dim：输入维数；out_dim：嵌入维数；hidden：0 表示线性层，> 0 表示一个隐层（ReLU）的 MLP； seed：初始化权重用的随机种子（None 表示不设种子——不推荐，结果每次运行都会不同） 【PyTorch 小知识】nn.Sequential 把几层按顺序串起来；nn.Linear(a, b) 是 y = xW^T + b，W 的形状是 (b, a)。 nn.Linear 创建时会用 PyTorch 的全局随机数生成器随机初始化权重。全局生成器在每个新进程里的初始状态 是随机的，所以**必须在创建模型之前**调用 torch.manual_seed，初始化才可复现。 """ def __init__(self, in_dim, out_dim=32, hidden=0, seed=0): super().__init__() if seed is not None: torch.manual_seed(seed) if hidden: self.net = nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU(), nn.Linear(hidden, out_dim)) else: self.net = nn.Linear(in_dim, out_dim) def forward(self, x): return self.net(x) class TextEncoder(nn.Module): """ 文字编码器：词袋（bag of words）。每个词有一个向量，一段文字的向量 = 它所有词向量的平均，再过一个线性层。 【PyTorch 小知识】nn.EmbeddingBag(V, k, mode="mean") 一次处理一批长短不一的文字： 把所有文字的词编号拼成一个长向量 ids，再给一个 offsets 告诉它每段文字从哪里开始， 它会对每一段分别取平均。比"补齐到同样长度再取平均"更省事。 【边界】一个词都没有的文字（空列表）会得到全 0 的向量。 """ def __init__(self, vocab_size, out_dim=32, emb_dim=32): super().__init__() self.emb = nn.EmbeddingBag(vocab_size, emb_dim, mode="mean") self.proj = nn.Linear(emb_dim, out_dim) def forward(self, ids, offsets): return self.proj(self.emb(ids, offsets)) class TinyCLIP(nn.Module): """ 迷你 CLIP = 图片编码器 + 文字编码器 + 一个可学习的温度。 【参数】in_dim：图片特征维数；vocab_size：词表大小；dim：共同的嵌入维数；hidden：图片编码器的隐层宽度（0 = 线性）； seed：初始化用的随机种子 """ def __init__(self, in_dim, vocab_size, dim=32, hidden=0, seed=0): super().__init__() self.image = ImageEncoder(in_dim, dim, hidden, seed=seed) self.text = TextEncoder(vocab_size, dim) self.log_scale = nn.Parameter(torch.tensor(float(np.log(1 / 0.07)))) def scale(self): """温度的倒数 τ = exp(log_scale)，截到 100 以内（CLIP 原文的做法），防止训练不稳定。""" return self.log_scale.exp().clamp(max=100.0) def encode_captions(captions, word_id): """ 把一批说明文字（token 列表）变成 EmbeddingBag 需要的 (ids, offsets)。 【参数】captions：token 列表的列表；word_id：词 → 编号的字典 【返回】(ids, offsets)：两个 LongTensor 【边界】不在词表里的词被忽略（在这个合成世界里不会发生）。 【Python 小知识】itertools 也能做，但这里用最朴素的循环，便于读懂： offsets[i] = 第 i 段文字的第一个词在 ids 里的位置。 """ ids, offsets = [], [] for cap in captions: offsets.append(len(ids)) ids.extend(word_id[w] for w in cap if w in word_id) return torch.tensor(ids, dtype=torch.long), torch.tensor(offsets, dtype=torch.long) # ============================================================================= # 训练：对比学习 # ============================================================================= def train_clip(model, x, captions, word_id, steps=1500, batch=256, lr=3e-3, seed=0): """ 用 InfoNCE 损失训练迷你 CLIP。 【参数】model：TinyCLIP；x：图片特征（numpy，形状 (n, d)）；captions：说明文字；word_id：词表； steps：训练步数；batch：每步的图文对数；lr：学习率（Adam）；seed：抽 batch 用的种子 【返回】训练好的 model（原地修改，同时返回它，方便链式调用） 【思路】每一步随机抽 batch 对，算 B × B 的相似度矩阵，对行、对列各做一次交叉熵（正确答案在对角线上）。 batch 里的其他 B − 1 段文字自动充当"负例"——这就是对比学习不需要标签的原因。 【边界】如果一个 batch 里有两张图的说明文字一模一样（例如都是 "husky photo"），对角线以外也有"正确答案"， InfoNCE 会把它当成负例——这是对比学习的一个已知缺陷，batch 越大越容易发生；这里的影响很小。 【PyTorch 小知识】F.cross_entropy(logits, target) 内部先做 log-softmax 再取负对数似然； torch.arange(B) 生成 [0, 1, ..., B−1]，正好是"第 i 行的正确答案是第 i 列"。 """ gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) X = torch.tensor(x, dtype=torch.float32) ids_all = [torch.tensor([word_id[w] for w in c if w in word_id], dtype=torch.long) for c in captions] opt = torch.optim.Adam(model.parameters(), lr=lr) n = len(X) target = torch.arange(min(batch, n)) for _ in range(steps): idx = torch.randperm(n, generator=gen)[:batch] lens = [len(ids_all[i]) for i in idx.tolist()] ids = torch.cat([ids_all[i] for i in idx.tolist()]) offsets = torch.tensor(np.concatenate([[0], np.cumsum(lens)[:-1]]), dtype=torch.long) img = F.normalize(model.image(X[idx]), dim=1) txt = F.normalize(model.text(ids, offsets), dim=1) logits = model.scale() * img @ txt.T loss = (F.cross_entropy(logits, target[:len(idx)]) + F.cross_entropy(logits.T, target[:len(idx)])) / 2 opt.zero_grad() loss.backward() opt.step() return model def zero_shot_head(model, class_names, word_id, prompt_words=(("photo",),)): """ 零样本分类头：把每个类别名写成提示，用文字编码器得到 C 个类别向量（归一化），对几种提示取平均。 【参数】model：TinyCLIP；class_names：类别名列表；word_id：词表； prompt_words：每种提示除类别名之外的词，例如 (("photo",), ("drawing",)) 表示两种提示 "husky photo" 和 "husky drawing"——多种提示取平均就是"提示集成"（prompt ensembling） 【返回】形状 (C, dim) 的 torch 张量（不需要梯度） """ with torch.no_grad(): embs = [] for words in prompt_words: caps = [[name] + list(words) for name in class_names] ids, offsets = encode_captions(caps, word_id) embs.append(F.normalize(model.text(ids, offsets), dim=1)) head = torch.stack(embs).mean(dim=0) return F.normalize(head, dim=1) def zero_shot_scores(model, head, x): """零样本分数：图片向量（归一化）和每个类别向量的余弦相似度。返回 numpy 数组 (n, C)。""" with torch.no_grad(): img = F.normalize(model.image(torch.tensor(x, dtype=torch.float32)), dim=1) return (img @ head.T).numpy() # ============================================================================= # 训练：监督学习、微调、参数插值 # ============================================================================= class Classifier(nn.Module): """ 图片编码器 + 线性分类头。分类头可以用零样本的类别向量初始化（这时它就是"零样本模型"本身）。 【参数】encoder：ImageEncoder；n_classes：类别数；dim：嵌入维数； init_head：None 表示随机初始化；否则是 (C, dim) 的张量，分类头权重设为 τ · 类别向量 normalize：True 时先把嵌入归一化（和 CLIP 的零样本分类一致） """ def __init__(self, encoder, n_classes, dim=32, init_head=None, scale=1.0, normalize=False, seed=0): super().__init__() self.encoder = encoder if seed is not None: torch.manual_seed(seed) self.head = nn.Linear(dim, n_classes) self.normalize = normalize if init_head is not None: with torch.no_grad(): self.head.weight.copy_(scale * init_head) self.head.bias.zero_() def forward(self, x): z = self.encoder(x) if self.normalize: z = F.normalize(z, dim=1) return self.head(z) def scores(self, x): """numpy 输入 → numpy 分数，方便和第 1 讲的 accuracy() 一起用。""" with torch.no_grad(): return self(torch.tensor(x, dtype=torch.float32)).numpy() def train_supervised(model, x, y, steps=1500, batch=256, lr=3e-3, seed=0): """ 用交叉熵训练一个 Classifier（编码器和分类头一起训练）。 【参数】model：Classifier；x：特征（numpy）；y：标签（numpy 整数数组）；其余同 train_clip 【返回】训练好的 model 【边界】y 里不能有 −1（调用方先去掉"硬匹配不上"的样本）。 """ gen = torch.Generator().manual_seed(seed) torch.manual_seed(seed) X = torch.tensor(x, dtype=torch.float32) Y = torch.tensor(y, dtype=torch.long) opt = torch.optim.Adam(model.parameters(), lr=lr) n = len(X) for _ in range(steps): idx = torch.randperm(n, generator=gen)[:batch] loss = F.cross_entropy(model(X[idx]), Y[idx]) opt.zero_grad() loss.backward() opt.step() return model def zero_shot_classifier(clip_model, head): """ 把零样本模型写成一个 Classifier：编码器 = CLIP 的图片编码器（拷贝），分类头 = τ · 类别向量。 【为什么】这样零样本模型、微调后的模型、插值后的模型就是同一种结构，可以直接做参数插值。 【返回】Classifier（normalize=True） """ enc = copy.deepcopy(clip_model.image) with torch.no_grad(): scale = float(clip_model.scale()) return Classifier(enc, head.shape[0], head.shape[1], init_head=head, scale=scale, normalize=True) def interpolate(model_a, model_b, alpha): """ WiSE-FT（Wortsman et al. 2022）：两个**同结构**模型的参数逐个做线性插值，θ = (1 − α)·θ_a + α·θ_b。 【参数】model_a：零样本模型；model_b：微调后的模型；alpha：0 = 完全是 a，1 = 完全是 b 【返回】一个新模型（不修改输入） 【为什么能行】微调是从零样本模型出发的，两者在参数空间里"连在一起"，中间的点通常仍然是好模型。 Wortsman 等人发现中间的点常常同时兼顾分布内和分布外。 """ out = copy.deepcopy(model_a) sa, sb = model_a.state_dict(), model_b.state_dict() out.load_state_dict({k: (1 - alpha) * sa[k] + alpha * sb[k] for k in sa}) return out # ============================================================================= # 自检 # ============================================================================= if __name__ == "__main__": import os import sys sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01")) from classifiers import accuracy from domain_world import DIM, DomainWorld world = DomainWorld(seed=0) web = world.web_pool(20000, np.random.default_rng(1)) clip = TinyCLIP(DIM, len(world.vocab)) train_clip(clip, web.x, web.captions, world.word_id, steps=1500) head = zero_shot_head(clip, world.class_names, world.word_id) suite = world.test_suite(50, np.random.default_rng(2)) for name, ts in suite.items(): print(f"零样本 {name:<10} {accuracy(zero_shot_scores(clip, head, ts.x), ts.label):.3f}")