""" qformer_lab.py —— Q-Former（BLIP-2）的缩小版：用几个可学习的查询向量，把图像交给冻结的语言模型（第 15 讲） ===================================================================== 【这个文件要回答的问题】 第 13 讲的 VLM（PaliGemma / LLaVA 式）用一个线性投影把图像特征直接变成语言模型的输入 token：图像有多少块， 就塞进去多少个 token（或者先平均池化到少一些）。语言模型要自己从这些 token 里找有用的信息。 BLIP-2（Li et al., 2023）的 Q-Former 换了一种思路： - 一组**可学习的查询向量**（BLIP-2 用 32 个），通过**交叉注意力**去"问"冻结的图像编码器（第 14 讲的交叉注意力： 查询来自这组向量，键和值来自图像块的特征）； - 不管图像有多少块，交给语言模型的永远只有这几个查询的输出（再经过一个线性投影）； - InstructBLIP（Dai et al., 2023）再进一步：把**问题**也送进 Q-Former，让查询"带着问题去看图"。 第 13 讲留下了一个问题：语言模型冻结时（阶段 A），"是不是某种形状"这类要把问题和图像放在一起比较的题很不稳定。 看问题的 Q-Former 能把比较挪到 Q-Former 里做吗？01 号脚本回答这个问题。 【本文件提供】 QFormerLayer / QFormer 查询（和问题）之间的自注意力 → 查询对图像块的交叉注意力 → 前馈层 LinearAdapter / QFormerAdapter 两种"图像 → k 个语言模型输入 token"的接口 AdapterVLM 接口 + 语言模型（对话格式同第 13 讲： [图像 token ×k] 问题 回答 ） setup 读取第 13 讲的 CLIP、第 11 讲的聊天助手，扩充词表，准备图像特征 make_batch / train_vlm / set_stage_a / answer / evaluate 【复用】第 13 讲 03 号脚本的问题生成（question、QTYPES、QNAMES），第 13 讲缓存的 CLIP（lecture13/checkpoints/clip_clip.pt）， 第 11 讲缓存的指令微调模型（lecture11/checkpoints/chat_sft.pt），第 14 讲的 CrossAttention。都只读。 """ import copy import importlib import os import random import sys import torch import torch.nn as nn import torch.nn.functional as F _HERE = os.path.dirname(os.path.abspath(__file__)) for _d in ("lecture14", "lecture13", "lecture12", "lecture11", "lecture02"): _p = os.path.join(_HERE, "..", _d) if _p not in sys.path: sys.path.insert(0, _p) from t2i_lab import CrossAttention # noqa: E402 from caption_world import HOLDOUT, HORIZ_WORDS, SHAPE_WORDS, SIZE_WORDS, VERT_WORDS, make_captioned # noqa: E402 from chat_lab import ASSISTANT, USER, Vocab, evaluate_tasks, expand_vocab, load_chat_sft # noqa: E402 from lm_models import FeedForward, MultiHeadAttention # noqa: E402 from vlm_lab import load_or_train_clip # noqa: E402 # 【Python 小知识】文件名以数字开头（03_vlm.py）不能写成 import 03_vlm，但 importlib.import_module("03_vlm") 可以 _vlm03 = importlib.import_module("03_vlm") QTYPES, QNAMES, question = _vlm03.QTYPES, _vlm03.QNAMES, _vlm03.question # ===================================================================== # 1. Q-Former # ===================================================================== class QFormerLayer(nn.Module): """ Q-Former 的一层（BLIP-2 的图 2，缩小版）： 1. 自注意力：查询之间互相看；如果送了问题，查询和问题的词拼成一个序列一起做（InstructBLIP）， 查询就能"读到"问题； 2. 交叉注意力：只有查询去看图像块的特征（键和值 = 冻结的图像编码器输出）——问题的词不直接看图； 3. 前馈层。每个子层前有 LayerNorm、后有残差（pre-norm）。 【参数】d：Q-Former 的宽度；d_img：图像特征的维度；heads：头数 """ def __init__(self, d, d_img, heads=4): super().__init__() self.ln1, self.ln2, self.ln3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d) self.attn = MultiHeadAttention(d, heads, causal=False) self.xattn = CrossAttention(d, d_img, heads) self.ffn = FeedForward(d, 4 * d) def forward(self, h, n_q, img, pad_mask=None): """【参数】h (B, n_q + L, d)：前 n_q 个是查询；img (B, 49, d_img)；pad_mask：(B, 1, 1, n_q + L)，0 或 −1e9""" self.attn.extra_mask = pad_mask h = h + self.attn(self.ln1(h)) self.attn.extra_mask = None q = h[:, :n_q] q = q + self.xattn(self.ln2(q), img, torch.ones(img.shape[:2], dtype=torch.bool)) h = torch.cat([q, h[:, n_q:]], 1) return h + self.ffn(self.ln3(h)) class QFormer(nn.Module): """ 【参数】n_q：查询个数（交给语言模型的图像 token 数）；d、n_layers；d_img；vocab_size：问题的词表大小（instruct 时用）； instruct：True 时把问题的词也送进来（InstructBLIP） 【输入】img (B, 49, d_img)；q_ids (B, L)：问题的编号（右边补 = 0），instruct=False 时忽略 【返回】(B, n_q, d)：每个查询的输出 【边界】查询的初始值 ~ N(0, 0.02²)，可学习；BLIP-2 的 Q-Former 从 BERT 初始化、有 1 亿参数，我们的只有几万。 """ def __init__(self, n_q=4, d=64, n_layers=2, d_img=64, vocab_size=0, instruct=False, max_q=12): super().__init__() self.n_q, self.instruct = n_q, instruct self.queries = nn.Parameter(torch.randn(n_q, d) * 0.02) if instruct: self.tok = nn.Embedding(vocab_size, d) self.pos = nn.Embedding(max_q, d) self.layers = nn.ModuleList(QFormerLayer(d, d_img) for _ in range(n_layers)) self.norm = nn.LayerNorm(d) def forward(self, img, q_ids=None): B = len(img) h = self.queries.expand(B, -1, -1) mask = None if self.instruct: txt = self.tok(q_ids) + self.pos(torch.arange(q_ids.size(1))) h = torch.cat([h, txt], 1) ok = torch.cat([torch.ones(B, self.n_q, dtype=torch.bool), q_ids != 0], 1) mask = ((~ok).float() * -1e9)[:, None, None, :] for layer in self.layers: h = layer(h, self.n_q, img, mask) return self.norm(h[:, :self.n_q]) # ===================================================================== # 2. 两种"图像 → k 个语言模型输入 token"的接口 # ===================================================================== def small_linear(d_in, d_out): """ 投影层：权重 ~ N(0, 0.02²)、偏置 0（和第 13 讲 03 号脚本一样）。 【为什么不用 PyTorch 的默认初始化】默认初始化的权重大得多（均匀分布 ±1/√d_in），图像 token 一开始就和词嵌入的尺度差很远； 我们试过：线性投影用默认初始化时，"是不是某种形状"三个种子都在 0.5 左右（见 README）。 """ lin = nn.Linear(d_in, d_out) nn.init.normal_(lin.weight, 0, 0.02) nn.init.zeros_(lin.bias) return lin class LinearAdapter(nn.Module): """ 第 13 讲的做法：一个线性投影。mode = "cls"（CLIP 的汇总向量，1 个 token）、"pool16"（7×7 的块平均池化成 4×4， 16 个 token）或 "all49"（49 个块全部送进去）。 """ def __init__(self, mode, d_img=64, d_lm=64): super().__init__() self.mode = mode self.k = {"cls": 1, "pool16": 16, "all49": 49}[mode] self.proj = small_linear(d_img, d_lm) def forward(self, cls, patches, q_ids): if self.mode == "cls": x = cls[:, None] elif self.mode == "pool16": grid = patches.transpose(1, 2).reshape(len(cls), -1, 7, 7) x = F.adaptive_avg_pool2d(grid, 4).flatten(2).transpose(1, 2) else: x = patches return self.proj(x) class QFormerAdapter(nn.Module): """Q-Former + 线性投影（BLIP-2 的"全连接层把查询的输出投到语言模型的输入维度"）。""" def __init__(self, n_q, vocab_size, instruct=False, d=64, n_layers=2, d_img=64, d_lm=64): super().__init__() self.k = n_q self.qformer = QFormer(n_q, d, n_layers, d_img, vocab_size, instruct) self.proj = small_linear(d, d_lm) def forward(self, cls, patches, q_ids): return self.proj(self.qformer(patches, q_ids)) class AdapterVLM(nn.Module): """ 接口 + 语言模型（第 2 讲结构的 GPT）。输入 = [ ] + k 个图像 token + [问题 … 回答 …]。 【返回】logits (N, 2 + k + len(after), V) """ def __init__(self, lm, adapter): super().__init__() self.lm, self.adapter = lm, adapter def forward(self, cls, patches, q_ids, ids_before, ids_after): x = torch.cat([self.lm.tok(ids_before), self.adapter(cls, patches, q_ids), self.lm.tok(ids_after)], 1) x = x + self.lm.pos(torch.arange(x.size(1))) for blk in self.lm.blocks: x = blk(x) return self.lm.head(self.lm.ln_f(x)) # ===================================================================== # 3. 准备：模型、词表、图像特征 # ===================================================================== def setup(quick=False, log=print): """ 【返回】字典：lm0（扩充了词表的聊天助手）、vocab、chat_vocab、sft（原来的聊天助手）、 tr / te：各自的 (cls (N, 64), patches (N, 49, 64), attrs (N, 4))——CLIP 图像编码器（冻结）的输出 """ clip, _ = load_or_train_clip("clip_clip", 30 if quick else 1500, loss="clip", log=log) for p in clip.parameters(): p.requires_grad_(False) sft, chat_vocab = load_chat_sft(log=log) new_words = [w for w in list(SHAPE_WORDS) + list(SIZE_WORDS) + list(VERT_WORDS) + list(HORIZ_WORDS) + ["yes", "no", "or", "where", "shape", "image", "big", "a", "at", "is", "it"] if w not in chat_vocab.stoi] new_words = list(dict.fromkeys(new_words)) vocab = Vocab(chat_vocab.itos[3:] + new_words) # 和第 13 讲 03 号脚本完全一样的词表 lm0 = expand_vocab(sft, chat_vocab, vocab) x_tr, a_tr, _ = make_captioned(8192, 1) keep = torch.tensor([tuple(a.tolist()) not in HOLDOUT for a in a_tr]) x_tr, a_tr = x_tr[keep], a_tr[keep] x_te, a_te, _ = make_captioned(500, 55) @torch.no_grad() def feats(x): out = [clip.img(x[i:i + 1000]) for i in range(0, len(x), 1000)] return torch.cat([o[0] for o in out]), torch.cat([o[1] for o in out]) return {"lm0": lm0, "vocab": vocab, "chat_vocab": chat_vocab, "sft": sft, "tr": feats(x_tr) + (a_tr,), "te": feats(x_te) + (a_te,)} # ===================================================================== # 4. 训练与评估 # ===================================================================== def pad_ids(seqs, value=0): T = max(len(s) for s in seqs) out = torch.full((len(seqs), T), value, dtype=torch.long) for i, s in enumerate(seqs): out[i, :len(s)] = torch.tensor(s) return out def make_batch(vocab, attrs, idx, rng, qtypes=QTYPES): """ 一批训练数据：每张图随机问一种问题。 【返回】(q_ids：问题本身（给看问题的 Q-Former）, ids_before, ids_after, targets_after) """ before = torch.tensor([[vocab.bos_id, vocab.stoi[USER]]] * len(idx)) qs, seqs = [], [] for i in idx: q, a = question(rng.choice(qtypes), attrs[i], rng) qs.append(vocab.encode(q)) seqs.append((vocab.encode(q + [ASSISTANT]), vocab.encode(a) + [vocab.eos_id])) after = pad_ids([q + a for q, a in seqs], vocab.pad_id) tgt = torch.full(after.shape, -100, dtype=torch.long) for b, (q, a) in enumerate(seqs): for j, w in enumerate(a): tgt[b, len(q) - 1 + j] = w # 预测回答的第一个词，以此类推 return pad_ids(qs, vocab.pad_id), before, after, tgt def set_stage_a(vlm): """ 阶段 A（第 13 讲）：语言模型的 Transformer 各层冻结，只训练接口（投影层 / Q-Former）和词表的两端（新加的词要能说出来）。 【返回】可训练的参数列表 """ for p in vlm.parameters(): p.requires_grad_(False) for m in (vlm.adapter, vlm.lm.tok, vlm.lm.head): for p in m.parameters(): p.requires_grad_(True) return [p for p in vlm.parameters() if p.requires_grad] def train_vlm(vlm, vocab, data, steps, params, lr=1e-3, batch=64, seed=0): """训练（只更新 params）。data = (cls, patches, attrs)。""" cls, patches, attrs = data rng = random.Random(seed) gen = torch.Generator().manual_seed(seed) opt = torch.optim.Adam(params, lr=lr) n_before = 2 + vlm.adapter.k vlm.train() for _ in range(steps): idx = torch.randint(len(attrs), (batch,), generator=gen).tolist() q_ids, before, after, tgt = make_batch(vocab, attrs, idx, rng) logits = vlm(cls[idx], patches[idx], q_ids, before, after)[:, n_before:] loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1), ignore_index=-100) opt.zero_grad() loss.backward() opt.step() vlm.eval() @torch.no_grad() def answer(vlm, vocab, cls, patches, questions, max_new=9): """贪心生成回答（同一批的问题长度相同）。【返回】词列表的列表""" n = len(questions) before = torch.tensor([[vocab.bos_id, vocab.stoi[USER]]] * n) q_ids = torch.tensor([vocab.encode(q) for q in questions]) after = torch.tensor([vocab.encode(q + [ASSISTANT]) for q in questions]) out, done = [[] for _ in range(n)], [False] * n for _ in range(max_new): nxt = vlm(cls, patches, q_ids, before, after)[:, -1].argmax(-1) for b, t in enumerate(nxt.tolist()): if not done[b]: if t == vocab.eos_id: done[b] = True else: out[b].append(vocab.itos[t]) if all(done): break after = torch.cat([after, nxt[:, None]], 1) return out def evaluate(vlm, vocab, data, seed=5): """每种问题问一遍所有测试图，逐词完全正确才算对。【返回】{问题类型: 准确率}""" cls, patches, attrs = data rng = random.Random(seed) acc = {} for qt in QTYPES: qa = [question(qt, attrs[i], rng) for i in range(len(attrs))] # 同一种问题的长度都相同，只有"是不是某种形状"里形状词不同——长度仍然相同 outs = answer(vlm, vocab, cls, patches, [q for q, _ in qa]) acc[qt] = sum(o == a for o, (_, a) in zip(outs, qa)) / len(qa) return acc def farm_accuracy(vlm, vocab): """第 11 讲 9 种农场任务（每种 30 题）的平均准确率：语言模型原来的本事还剩多少。""" return sum(evaluate_tasks(vlm.lm, vocab, 30, 7, "train")[0].values()) / 9 __all__ = ["QTYPES", "QNAMES", "question", "small_linear", "QFormerLayer", "QFormer", "LinearAdapter", "QFormerAdapter", "AdapterVLM", "setup", "make_batch", "set_stage_a", "train_vlm", "answer", "evaluate", "farm_accuracy"]