ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch实现原型网络:从少样本学习到度量学习的完整实践

PyTorch实现原型网络:从少样本学习到度量学习的完整实践

简介:一份使用PyTorch实现的原型网络(Prototypical Networks)少样本学习代码库,面向机器学习开发者、算法研究人员以及正在学习元学习与少样本分类的进阶人群,可用于快速复现基于该经典模型的实验流程。压缩包共12个文件,以py源文件为主体,涵盖模型定义、训练脚本、原型损失函数、批次采样器、Omniglot数据集加载与参数解析等模块;另含png示意图、README说明、MIT许可证及gitignore配置。代码结构清晰,便于直接迁移到自定义数据与实验任务中。整体仅135KB,属于轻量级参考实现。目前已有920人学习下载,适合希望对照论文公式理解并动手实践原型的读者。通过这份代码,可掌握prototypical networks的核心思想、距离度量与损失计算方式,并基于PyTorch完成数据加载、训练与可视化验证。

1. Prototypical Networks:Few-shot Learning 里最不需要调参的原型方法

Few-shot Learning 的概念听上去吓人,但 Prototypical Networks 的核心逻辑一句话就能说清:每个类别用它的样本在嵌入空间里的均值作为“类原型”,查询样本离哪个原型最近,就判给哪个类。你不需要微调一个庞大的预训练模型,也不需要设计复杂的元学习更新规则,把编码器训好,剩下全是距离计算。这个项目用 PyTorch 完整实现了这一流程,从 episode 采样、原型计算到损失函数和评估循环全部可跑,适合刚接触小样本学习、想复现论文基线,或者要在自己数据集上快速验证想法的人。

我拆这个项目时最深的感受是:它把“少样本”问题变成了一个度量学习问题,而度量学习一旦嵌入空间质量过关,分类效果就水到渠成。下面从原理、数据管线、模型实现到踩坑记录,按我实际复现的顺序展开。

2. 算法拆解:类原型、欧氏距离与 episode 训练机制

2.1 从 N-way K-shot 说起,Few-shot 的默认设定

Few-shot 里最常用的评测协议是 N-way K-shot。N 表示每个 episode 里出现的类别数,K 表示每个类别给多少张支撑集样本。比如 5-way 1-shot,就是每个 episode 随机抽 5 个类,每类只给 1 张 labeled 样本,剩下同类样本作为查询集用来计算准确率。这个设定之所以残酷,是因为模型在每个 episode 里看到的可用标注数据极少,必须依靠嵌入空间里“同类聚拢、异类分开”的分布特性来完成分类。

支撑集和查询集的分工是硬性的:支撑集样本用来计算类原型,查询集样本用来测试模型能不能正确归类。训练阶段每个 episode 都重新采样支撑集和查询集,相当于不断给模型出新的“限时归类题”,逼迫编码器学会通用特征,而不是记住某个类别。这也是原型网络和普通分类网络最大的差异——没有固定的全连接分类头,分类行为完全在嵌入空间里动态发生。

2.2 类原型定义与欧氏距离的选用理由

原型网络的定义形式很直接:对每个类别 c,取其支撑集样本经过编码器 f_φ 映射后的嵌入向量,求均值作为原型 p_c:

p_c = (1 / |S_c|) * Σ f_φ(x_i) 其中 x_i ∈ S_c

查询样本 x 的类别概率通过 softmax 计算,logits 取负的平方欧氏距离:

p(y = c | x) = exp(-d(f_φ(x), p_c)) / Σ_c' exp(-d(f_φ(x), p_c'))

这里 d 用平方欧氏距离,不是余弦相似度。论文里有一条关键推导:当嵌入函数是线性时,平方欧氏距离等价于高斯密度估计下的对数似然,也就是说原型网络可以看成在嵌入空间里对每个类做高斯均值估计,方差相等时最大似然解恰好就是算术平均。这解释了为什么用均值做原型不是拍脑袋,而是有概率解释的。

距离/相似度适用场景注意点
平方欧氏距离原型网络论文原始设定,稳定logits = -d²,距离越大得分越低
余弦相似度嵌入向量做了 L2 归一化需要额外归一化层,原论文不采用
曼哈顿距离极少见对维度尺度敏感,不推荐

我自己的经验是,尽量守住欧氏距离这个原始配置,不要一上来就换相似度度量。模型不收敛时换度量会引入额外变量,排错更困难。

2.3 episode 训练机制:为什么每个 batch 是一局游戏

传统的分类训练每个 batch 从整个数据集里随机抽样本,类别分布天然均衡即可。原型网络不一样,它每个 batch 就是一个完整的 episode:先随机选 N 个类,每类抽 K 张做支撑集,再每类抽 Q 张做查询集,然后用支撑集算原型、对查询集算交叉熵损失。

这种机制让训练和评测的分布完全对齐。模型在训练时经历的就是“只给几张图,马上要分类”的场景,不会像迁移学习那样先把类别学全再硬套少样本任务。PyTorch 里实现这个逻辑需要自定义 Dataset 和 Sampler,下面第 3 章我会给出可直接用的版本。

3. 环境准备与数据管线:先把 episode 采样器跑通

3.1 PyTorch 环境搭建与项目结构

复现这个项目之前,先把环境准备好。我习惯用 conda 管理 Python 环境,PyTorch 的安装版本取决于你的 CUDA 版本。如果机器上有 NVIDIA 显卡且驱动支持 CUDA 12.1,直接装 GPU 版本;否则先用 CPU 版本跑通逻辑,等有 GPU 机器再切换。

conda create -n proto_net python=3.9 -y conda activate proto_net pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 pip install numpy scikit-learn matplotlib tqdm

这段命令里,torch和torchvision用 cu121 的 wheel 源,会匹配 CUDA 12.1 的运行时;如果安装后运行torch.cuda.is_available()返回 False,说明版本不匹配,可以换成cpu的 index-url 再装。项目本身不依赖额外重型库,标准目录分三块:data/放数据集,models/放编码器定义,train.py和eval.py分别处理训练与评测。

提示:conda 无法识别为命令时,先检查 Anaconda 是否加入了 PATH,或者改用source activate的方式。

3.2 支撑集与查询集的采样实现

自定义采样器是数据管线的核心。下面这段代码实现了一个 episode 级别的 BatchSampler:每次从类别列表里随机抽取 N 个类,再从每个类中选 K + Q 个样本,前 K 个当支撑集,后 Q 个当查询集。为了防止两张同类图在同一个 episode 里被既当支撑又当查询,我用索引偏移做了严格区隔。

import numpy as np from torch.utils.data import Sampler class EpisodeBatchSampler(Sampler): def __init__(self, labels, n_way, k_shot, n_query, episodes): self.labels = np.array(labels) self.classes = np.unique(self.labels) self.n_way = n_way self.k_shot = k_shot self.n_query = n_query self.episodes = episodes self.class_to_indices = { cls: np.where(self.labels == cls)[0] for cls in self.classes } def __iter__(self): for _ in range(self.episodes): chosen_classes = np.random.choice(self.classes, self.n_way, replace=False) episode_indices = [] for cls in chosen_classes: pool = self.class_to_indices[cls] selected = np.random.choice(pool, self.k_shot + self.n_query, replace=False) support_idx = selected[:self.k_shot] query_idx = selected[self.k_shot:] episode_indices.extend(support_idx.tolist()) episode_indices.extend(query_idx.tolist()) yield episode_indices def __len__(self): return self.episodes

逻辑说明:外层循环的每轮迭代生成一个完整 episode 的样本索引列表。每个类别内部先一次性抽取 K + Q 个不重复索引,再按位置切分,确保支撑集和查询集没有交集。replace=False是关键,如果某个类别的样本数少于 K + Q,这里会直接报错,所以数据预处理阶段必须检查每个类至少有多少张图。

参数说明:n_way设为 5 就是 5-way 分类,k_shot设为 1 就是 1-shot,n_query通常设为 15,太小会导致每个 episode 的准确率方差过大,评估不稳定。episodes控制迭代轮数,训练时我一般用 1000 个 episode 起步。

3.3 图像预处理与数据集适配

miniImageNet 和 Omniglot 这两个经典数据集的预处理方式有差异。miniImageNet 原始图像是 84×84×3 的 RGB 图,Omniglot 是 28×28 的单通道字符图。PyTorch 里可以用torchvision.transforms链式处理,但注意不要做过度增强——少样本场景下随机裁剪或颜色抖动幅度太大会让本就稀少的支撑样本失真。

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(84), transforms.ColorJitter(brightness=0.4, contrast=0.4), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) eval_transform = transforms.Compose([ transforms.Resize(84), transforms.CenterCrop(84), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

逻辑说明:训练和评测的 transform 必须分开。训练时用RandomResizedCrop做随机裁剪,相当于给模型一定的数据增强;评测时只能CenterCrop,保持输入固定,否则同一张图每次评估结果都不一样,无法复现指标。Normalize 用的均值方差来自 ImageNet,因为 miniImageNet 是从 ImageNet 采样的子集,沿用这套参数是最稳妥的做法。

参数说明:RandomResizedCrop(84)的裁剪区域随机,这样做对 1-shot 场景往往比固定缩放更有效,因为它强制模型关注目标的局部判别区域。ColorJitter 的幅度值我一般控制在 0.4 以内,过大会让字符类数据集的笔画特征被破坏。

4. 模型实现与训练循环:四层卷积编码器加原型损失

4.1 编码器选型:四层卷积为什么是默认配置

原型网络论文里使用的编码器是一个四层卷积网络,每层 64 个卷积核,后接 BatchNorm 和 ReLU,再跟一个 2×2 的 max pooling。这个结构放到今天看不算深,但在少样本任务里它有一个明显的优点:参数少,不容易过拟合,而且嵌入维度适中。

import torch import torch.nn as nn import torch.nn.functional as F class ConvEncoder(nn.Module): def __init__(self, in_channels=3, hidden_dim=64): super().__init__() self.features = nn.Sequential( nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), ) self.flatten_dim = hidden_dim * 5 * 5 # 84x84 -> 5x5 def forward(self, x): x = self.features(x) return x.view(x.size(0), -1)

逻辑说明:输入 84×84 的图经过四轮“卷积 + 池化”,空间尺寸从 84 依次降为 42、21、10、5,最后一层输出的特征图是 5×5×64,展平后得到 1600 维的嵌入向量。如果你用的是 Omniglot 的 28×28 输入,这个展平维度就要改成 64,因为四次池化后 28 会变成 1,输出是 1×1×64。

参数说明:hidden_dim=64是全篇出现最多的超参数,它同时控制嵌入维度。如果你把它改成 128,模型容量会增加,但训练速度和显存占用也水涨船高。我实测过 5-way 1-shot 场景,64 和 128 的准确率差距通常在 1 到 2 个点以内,优先用 64 做基线再去调。

4.2 原型计算与损失函数:先求均值,再算欧氏距离

模型前向传播拿到支撑集和查询集的嵌入向量后,需要一个函数计算原型并输出分类 logits。这里要注意顺序:先把支撑集按类别分组求均值,得到 N 个原型向量,再算查询集每个样本与 N 个原型的距离。

def prototypical_loss(prototypes, query_embeddings, query_labels, n_way): # prototypes: [n_way, d] # query_embeddings: [n_query * n_way, d] # query_labels: [n_query * n_way] dists = torch.cdist(query_embeddings, prototypes, p=2).pow(2) logits = -dists loss = F.cross_entropy(logits, query_labels) acc = (logits.argmax(dim=1) == query_labels).float().mean() return loss, acc

逻辑说明:torch.cdist一次算出查询集所有样本与所有原型的欧氏距离矩阵,矩阵形状是[查询集总数, n_way],每个元素表示一个查询样本到某个原型的距离。注意我对距离做了平方操作.pow(2),这一步与论文里平方欧氏距离的设定对应,如果忘了平方,实验结果的绝对值会整体变低。

参数说明:p=2是欧氏距离的 L2 范数。cross_entropy接收的 logits 是负距离,越接近 0 表示离原型越近,模型越确信该样本属于这个类。这里有个隐藏细节——查询集样本的数量等于n_query * n_way,所以 query_labels 的构造顺序必须和 episode 采样时一致,否则损失计算会错位。

4.3 训练循环:一个 episode 一个 batch 的 PyTorch 写法

有了采样器、编码器和损失函数,训练主循环就顺理成章了。下面是一个完整的训练流程,包含支撑集/查询集的拆分和优化器配置。

import torch.optim as optim def train_one_episode(model, support_x, support_y, query_x, query_y, optimizer): model.train() support_emb = model(support_x) query_emb = model(query_x) unique_labels = torch.unique(support_y) prototypes = [] for label in unique_labels: mask = support_y == label proto = support_emb[mask].mean(dim=0) prototypes.append(proto) prototypes = torch.stack(prototypes) loss, acc = prototypical_loss(prototypes, query_emb, query_y, len(unique_labels)) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item(), acc.item()

逻辑说明:支撑集和查询集分别过一次编码器得到两组嵌入向量,再按支撑集标签用布尔掩码选出同一类的嵌入,对它们取均值得到原型。torch.unique(support_y)的顺序就是原型的类别顺序,后续计算查询集准确率时依赖这个顺序的一致性,所以不要手动打乱 unique_labels。

参数说明:优化器我固定用 Adam,学习率从 1e-3 起步。原型网络对学习率相对不敏感,但如果 loss 在十几个 episode 后仍然没有下降趋势,把学习率降到 3e-4 再试。这里不需要学习率衰减策略,少样本任务训练周期短,余弦退火带来的收益不明显。

训练时的 episode 总轮数取决于数据集规模。miniImageNet 上我一般跑 10000 个 episode,每个 episode 用小批量优化一次,差不多 20 分钟内能看到收敛趋势;Omniglot 更小,3000 个 episode 就足够。

5. 避坑与常见问题:五个最影响复现结果的细节

5.1 类别泄漏:训练和评估用了同一批类别

现象:训练准确率涨得很快,评估准确率却比随机高不了多少;或者调换数据集划分后指标崩盘。

原因:episode 采样时类别池没有按数据集划分隔离。如果训练期的类别池包含了验证集和测试集的类别,模型在训练时已经见过这些类的样本,评估时相当于开卷考试,指标虚高。

解决:严格按数据集划分维护三个类别列表。训练采样器只能从训练类里采样,验证和测试采样器各自使用独立类别池。代码里self.classes = np.unique(self.labels)这句需要根据当前阶段传不同的 labels 数组进来,三段数据各跑一遍。

5.2 随机种子不固定,每次实验结果波动明显

现象:同一个配置跑两遍,准确率差了 3 个百分点以上,很难判断改动是否有效。

原因:episode 采样、数据增强、模型初始化都有随机性,少样本评估本身方差就大,不定种子等于所有对比都在噪声里做。

解决:在所有随机入口固定种子。

def set_seed(seed=42): import random random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)

我习惯把评估过程跑满 1000 个 episode,再报告平均准确率和 95% 置信区间。单个 episode 的结果没有参考价值,5-way 1-shot 的准确率波动可以到 ±5%。

5.3 距离矩阵里的 NaN:嵌入向量的值域失控

现象:训练几轮后 loss 变成 NaN,模型参数全部失效。

原因:BatchNorm 在 batch size 很小(支撑集总共就 5 个样本)时统计量不稳定,加上平方距离的数值范围大,梯度很容易爆炸。

解决:先把学习率降到 1e-4 试跑,排除优化器问题;然后在每个卷积层的 BatchNorm 之后确认前向输出方差在 1 附近。还有一个实用招数:把原型和查询嵌入都做一次 L2 归一化,虽然改了距离语义,但数值稳定性好很多。我是先做梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 10),再考虑动距离公式。

5.4 支撑集和查询集索引串位,标签顺序对不上

现象:代码不报错,但准确率一直徘徊在 1/N 附近,像没学一样。

原因:episode 采样器输出的索引顺序是“每个类先支撑后查询”,但 DataLoader 拿到索引后如果没按顺序拆分,支撑集里混入了查询集样本,原型均值被污染。

解决:DataLoader 必须配合我第 3 章的 EpisodeBatchSampler 使用,并且 collate 之后按位置切割:

split = n_way * k_shot support_x, query_x = batch_x[:split], batch_x[split:] support_y, query_y = batch_y[:split], batch_y[split:]

切割位置务必和采样器生成顺序完全一致,这是整个项目里最容易翻车的静默 bug。

5.5 每类样本数不足,采样器直接崩溃

现象:运行到一半报ValueError: Cannot take a larger sample than population。

原因:某个类别在数据集里的样本总数少于k_shot + n_query,比如某些类只有 3 张图,而配置要求 1-shot + 15 query。

解决:预处理阶段过滤掉样本量不足的类别,或者降低 n_query。我在做 miniImageNet 时发现原版数据集有些类只有 20 张图,n_query 设到 15 就是极限;Omniglot 每个类有 20 个手写样本,相对宽裕。

6. 进阶技巧:评估置信区间与嵌入空间可视化验证

把训练跑通只是第一步,真正判断模型好坏要看两件事:准确率是否稳定,以及嵌入空间是否真的按类别聚拢。先说评估端的置信区间计算。由于每个 episode 的准确率波动大,我习惯在所有测试 episode 结束后统一计算均值和 95% 置信区间,而不是只报一个点估计。

import numpy as np def evaluate(model, dataloader, n_way, k_shot, n_query): model.eval() accuracies = [] with torch.no_grad(): for episode in dataloader: support_x, support_y, query_x, query_y = episode support_emb = model(support_x) query_emb = model(query_x) prototypes = torch.stack([ support_emb[support_y == label].mean(dim=0) for label in torch.unique(support_y) ]) dists = torch.cdist(query_emb, prototypes, p=2).pow(2) preds = dists.argmin(dim=1) acc = (preds == query_y).float().mean().item() accuracies.append(acc) mean_acc = np.mean(accuracies) ci95 = 1.96 * np.std(accuracies) / np.sqrt(len(accuracies)) return mean_acc, ci95

1.96 * std / sqrt(n)算出来的是标准误差倍乘的置信半径,报告时写成mean_acc ± ci95就够专业。1000 个 episode 跑完,5-way 1-shot 的置信区间一般能压到 ±1.5% 以内,如果区间宽度超过 ±3%,说明评估轮数太少或测试类本身分布不均衡,需要加大 episode 数。

再说嵌入空间可视化。我会随机采样部分支撑集和查询集样本,把它们的嵌入向量丢进 t-SNE 做降维,然后按类别着色画散点图。这一步能直观看出训练是否欠拟合:如果 5 个类别的点完全交叠,说明编码器没有学到判别特征;如果每个类聚成独立团簇且查询点落在对应原型附近,说明模型真的理解了。

from sklearn.manifold import TSNE import matplotlib.pyplot as plt embeddings = torch.cat([support_emb, query_emb], dim=0).cpu().numpy() tsne = TSNE(n_components=2, perplexity=30, random_state=42).fit_transform(embeddings) colors = ['red', 'blue', 'green', 'orange', 'purple'] for i, label in enumerate(torch.unique(support_y)): plt.scatter(tsne[support_y == label, 0], tsne[support_y == label, 1], c=colors[i], marker='o', label=f'class {label.item()}') plt.legend() plt.show()

perplexity=30是 t-SNE 的默认推荐值,点数太少时降到 10 到 15 更合理。如果聚簇轮廓模糊但准确率尚可,通常是嵌入维度偏高导致低维投影失真,这不算模型问题,而是可视化手段的边界。

从那以后,我每次换数据集复现原型网络,都强制走一遍完整流程:先查每个类别的样本数量下限,再固定随机种子跑 100 个 episode 验证能收敛,最后才把评估轮数拉满到 1000。这套习惯帮我省掉了大量“实验结果忽高忽低”的无效对比时间。希望这篇拆解能让你复制项目时少走几步弯路。

本文还有配套的精品资源,点击获取

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进