"""第 40 章：从像素和两个词开始，训练一个小型图文检索模型。""" from argparse import ArgumentParser from copy import deepcopy from hashlib import sha256 from pathlib import Path import json import platform import numpy as np import torch from torch import nn from torch.nn import functional as F COLORS = ["red", "green", "blue"] SHAPES = ["circle", "square", "triangle"] CAPTIONS = [f"{color} {shape}" for color in COLORS for shape in SHAPES] RGB = np.array([[.85, .18, .16], [.15, .68, .30], [.15, .35, .88]]) HELD_OUT = [2, 7] # red triangle、blue square 不进入训练或验证。 SEEN = [i for i in range(9) if i not in HELD_OUT] SIZE, PATCH, WIDTH, DIM = 32, 8, 32, 16 SEED, EPOCHS, TEMPERATURE = 2026, 40, .10 REPORT = Path(__file__).with_name("data") / "cross-modal-alignment.json" def make_data(split): counts = {"train": 48, "validation": 16, "test": 24} seeds = {"train": 101, "validation": 202, "test": 303} rng = np.random.default_rng(seeds[split]) labels = range(9) if split == "test" else SEEN yy, xx = np.mgrid[:SIZE, :SIZE] images, targets, ids = [], [], [] for label in labels: for index in range(counts[split]): cx, cy = rng.uniform(12, 20, 2) radius, angle = rng.uniform(6, 10), rng.uniform(-.3, .3) dx, dy = xx - cx, yy - cy u = np.cos(angle) * dx + np.sin(angle) * dy v = -np.sin(angle) * dx + np.cos(angle) * dy if label % 3 == 0: mask = u*u + v*v <= radius*radius elif label % 3 == 1: mask = np.maximum(abs(u), abs(v)) <= radius * .88 else: mask = (v >= -radius) & (v <= radius) & (abs(u) <= (v+radius)/2) background = rng.uniform(.86, .99) color = np.clip(RGB[label//3] + rng.normal(0, .025, 3), 0, 1) image = np.full((SIZE, SIZE, 3), background) image[mask] = color image = np.clip(image + rng.normal(0, .01, image.shape), 0, 1) images.append(image.transpose(2, 0, 1).astype("float32")) targets.append(label) ids.append(f"{split}-{label}-{index:02d}") return torch.from_numpy(np.stack(images)), torch.tensor(targets), ids def text_ids(labels): # 六个词；词表只表达颜色与形状，不支持自由句子。 return torch.stack([labels // 3, 3 + labels % 3], dim=-1) class TinyImageEncoder(nn.Module): def __init__(self): super().__init__() self.patch = nn.Conv2d(3, WIDTH, PATCH, stride=PATCH) self.cls = nn.Parameter(torch.randn(1, 1, WIDTH) * .02) self.position = nn.Parameter(torch.randn(1, 17, WIDTH) * .02) self.blocks = nn.ModuleList([ nn.TransformerEncoderLayer(WIDTH, 4, WIDTH*2, dropout=0, activation="gelu", batch_first=True, norm_first=True) for _ in range(2) ]) self.norm = nn.LayerNorm(WIDTH) self.projection = nn.Linear(WIDTH, DIM, bias=False) def forward(self, images): x = self.patch(images * 2 - 1).flatten(2).transpose(1, 2) x = torch.cat([self.cls.expand(len(x), -1, -1), x], dim=1) x = x + self.position for block in self.blocks: x = block(x) # 图像可以同时看全部块，没有语言生成的因果掩码。 return self.projection(self.norm(x[:, 0])) class DualEncoder(nn.Module): def __init__(self): super().__init__() self.image = TinyImageEncoder() self.words = nn.Embedding(6, WIDTH) self.text_projection = nn.Linear(WIDTH, DIM, bias=False) def encode_image(self, images): return F.normalize(self.image(images), dim=-1) def encode_text(self, ids): return F.normalize(self.text_projection(self.words(ids).mean(1)), dim=-1) def forward(self, images, ids): return self.encode_image(images) @ self.encode_text(ids).T / TEMPERATURE def contrastive_loss(logits): if logits.ndim != 2 or logits.shape[0] != logits.shape[1] or len(logits) < 2: raise ValueError("需要至少两个一一配对的样本，得分矩阵必须为方阵") target = torch.arange(len(logits), device=logits.device) return (F.cross_entropy(logits, target) + F.cross_entropy(logits.T, target)) / 2 @torch.inference_mode() def scores(model, images, labels): model.eval() vectors = torch.cat([model.encode_image(x) for x in images.split(64)]) return vectors @ model.encode_text(text_ids(labels)).T def retrieval_metrics(similarity, image_labels, caption_labels): # 多张图可与同一句描述相关；按语义标签判断，不能只认矩阵对角线。 i_order = similarity.argsort(dim=1, descending=True, stable=True) t_order = similarity.T.argsort(dim=1, descending=True, stable=True) result = {} for name, order, gallery, query in [ ("i2t", i_order, caption_labels, image_labels), ("t2i", t_order, image_labels, caption_labels), ]: relevant = gallery[order] == query[:, None] for k in [1, 3]: result[f"{name}_hit{k}"] = relevant[:, :k].any(1).float().mean().item() return result def evaluate(model, data): images, labels, ids = data captions = torch.arange(9) similarity = scores(model, images, captions) prediction = similarity.argmax(1) known = torch.tensor([int(x) in SEEN for x in labels]) result = retrieval_metrics(similarity, labels, captions) result.update({ "seen_i2t_hit1": (prediction[known] == labels[known]).float().mean().item(), "heldout_i2t_hit1": (prediction[~known] == labels[~known]).float().mean().item(), "per_class_correct": [int(((prediction == label) & (labels == label)).sum()) for label in range(9)], "ids": ids, "labels": labels.tolist(), "predictions": prediction.tolist(), "similarities": similarity.tolist(), }) return result def tensor_digest(tensor): return sha256(tensor.contiguous().numpy().tobytes()).hexdigest() def search_images(model, query, images, k=3): if query not in CAPTIONS: raise ValueError("本例只接受九种颜色与形状组合，例如 blue square") similarity = scores(model, images, torch.tensor([CAPTIONS.index(query)]))[:, 0] order = similarity.argsort(descending=True, stable=True)[:k] return [(int(index), float(similarity[index])) for index in order] def run(): torch.set_num_threads(1) torch.manual_seed(SEED) torch.use_deterministic_algorithms(True) # 关闭推理专用 fast path，使训练与评估使用同一注意力实现路径。 torch.backends.mha.set_fastpath_enabled(False) train, validation, test = [make_data(s) for s in ["train", "validation", "test"]] model = DualEncoder() initial = evaluate(model, test) optimizer = torch.optim.AdamW(model.parameters(), lr=.001, weight_decay=.01) generator = torch.Generator().manual_seed(SEED + 1) image_bank = train[0].reshape(7, 48, 3, SIZE, SIZE) labels = torch.tensor(SEEN) best, best_loss, best_epoch = None, float("inf"), 0 history = [] for epoch in range(1, EPOCHS + 1): model.train() orders = torch.stack([torch.randperm(48, generator=generator) for _ in SEEN]) total = 0. for step in range(48): # 每批七对，每种语义只出现一次，避免同义描述被当成负样本。 images = image_bank[torch.arange(7), orders[:, step]] loss = contrastive_loss(model(images, text_ids(labels))) optimizer.zero_grad() loss.backward() optimizer.step() total += loss.item() if epoch == 1 or epoch % 5 == 0: val_logits = scores(model, validation[0], labels) / TEMPERATURE val_targets = torch.searchsorted(labels, validation[1]) val_loss = F.cross_entropy(val_logits, val_targets).item() history.append({"epoch": epoch, "train_loss": total/48, "validation_i2t_ce": val_loss}) if val_loss < best_loss: best, best_loss, best_epoch = deepcopy(model.state_dict()), val_loss, epoch model.load_state_dict(best) final = evaluate(model, test) state_hash = sha256() for name, tensor in sorted(best.items()): state_hash.update(name.encode()) state_hash.update(tensor.contiguous().numpy().tobytes()) report = { "config": {"seed": SEED, "epochs": EPOCHS, "temperature": TEMPERATURE, "size": SIZE, "patch": PATCH, "width": WIDTH, "dimension": DIM, "batch_size": 7, "learning_rate": .001, "weight_decay": .01, "held_out": HELD_OUT, "selection": "minimum seen validation i2t CE"}, "environment": {"python": platform.python_version(), "torch": torch.__version__, "numpy": np.__version__, "device": "cpu", "threads": 1}, "source_sha256": sha256(Path(__file__).read_bytes()).hexdigest(), "data_sha256": {name: tensor_digest(data[0]) for name, data in zip(["train", "validation", "test"], [train, validation, test])}, "model_sha256": state_hash.hexdigest(), "parameter_count": sum(p.numel() for p in model.parameters()), "selected_epoch": best_epoch, "history": history, "initial": initial, "trained": final, } return model, report def print_summary(report): print("合成图文：训练 336 张，验证 112 张，测试 216 张；候选描述 9 条。") print("留出组合：red triangle、blue square；训练与验证均不使用。") print(f"参数 {report['parameter_count']}；按验证损失选择第 {report['selected_epoch']} 轮。") for key, name in [("initial", "随机初始化"), ("trained", "对比训练后")]: row = report[key] print(f"{name}：图找文 Hit@1={row['i2t_hit1']:.3f}；" f"文找图 Hit@1={row['t2i_hit1']:.3f}；Hit@3={row['t2i_hit3']:.3f}") print(f" 已见组合={row['seen_i2t_hit1']:.3f}；" f"未见组合={row['heldout_i2t_hit1']:.3f}") print("每类图找文正确数（各 24 张）：", report["trained"]["per_class_correct"]) print("这些结果只描述合成图库；模型不支持自由句子，也不会生成回答。") def main(): parser = ArgumentParser(description=__doc__) parser.add_argument("--out", type=Path, help="另存本次训练的完整报告") parser.add_argument("--checkpoint", type=Path, help="另存模型权重") parser.add_argument("--query", choices=CAPTIONS, help="训练后用两个词检索测试图库") args = parser.parse_args() model, report = run() if args.out: args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n") if args.checkpoint: args.checkpoint.parent.mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), args.checkpoint) restored = DualEncoder() restored.load_state_dict(torch.load(args.checkpoint, weights_only=True)) torch.testing.assert_close(scores(model, make_data("test")[0], torch.arange(9)), scores(restored, make_data("test")[0], torch.arange(9)), rtol=0, atol=0) print_summary(report) if args.query: images, labels, ids = make_data("test") print(f"查询：{args.query}") for index, similarity in search_images(model, args.query, images): print(f" {ids[index]}；实际 {CAPTIONS[int(labels[index])]}；余弦 {similarity:.4f}") if __name__ == "__main__": main()