简介一份使用PyTorch实现的原型网络Prototypical Networks少样本学习代码库面向机器学习开发者、算法研究人员以及正在学习元学习与少样本分类的进阶人群可用于快速复现基于该经典模型的实验流程。压缩包共12个文件以py源文件为主体涵盖模型定义、训练脚本、原型损失函数、批次采样器、Omniglot数据集加载与参数解析等模块另含png示意图、README说明、MIT许可证及gitignore配置。代码结构清晰便于直接迁移到自定义数据与实验任务中。整体仅135KB属于轻量级参考实现。目前已有920人学习下载适合希望对照论文公式理解并动手实践原型的读者。通过这份代码可掌握prototypical networks的核心思想、距离度量与损失计算方式并基于PyTorch完成数据加载、训练与可视化验证。1. Prototypical NetworksFew-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_cp_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 python3.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, replaceFalse) episode_indices [] for cls in chosen_classes: pool self.class_to_indices[cls] selected np.random.choice(pool, self.k_shot self.n_query, replaceFalse) 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 个不重复索引再按位置切分确保支撑集和查询集没有交集。replaceFalse是关键如果某个类别的样本数少于 K Q这里会直接报错所以数据预处理阶段必须检查每个类至少有多少张图。参数说明n_way设为 5 就是 5-way 分类k_shot设为 1 就是 1-shotn_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(brightness0.4, contrast0.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_channels3, hidden_dim64): super().__init__() self.features nn.Sequential( nn.Conv2d(in_channels, hidden_dim, kernel_size3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, kernel_size3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, kernel_size3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, kernel_size3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), 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_dim64是全篇出现最多的超参数它同时控制嵌入维度。如果你把它改成 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, p2).pow(2) logits -dists loss F.cross_entropy(logits, query_labels) acc (logits.argmax(dim1) query_labels).float().mean() return loss, acc逻辑说明torch.cdist一次算出查询集所有样本与所有原型的欧氏距离矩阵矩阵形状是[查询集总数, n_way]每个元素表示一个查询样本到某个原型的距离。注意我对距离做了平方操作.pow(2)这一步与论文里平方欧氏距离的设定对应如果忘了平方实验结果的绝对值会整体变低。参数说明p2是欧氏距离的 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(dim0) 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(seed42): 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(dim0) for label in torch.unique(support_y) ]) dists torch.cdist(query_emb, prototypes, p2).pow(2) preds dists.argmin(dim1) 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, ci951.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], dim0).cpu().numpy() tsne TSNE(n_components2, perplexity30, random_state42).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], ccolors[i], markero, labelfclass {label.item()}) plt.legend() plt.show()perplexity30是 t-SNE 的默认推荐值点数太少时降到 10 到 15 更合理。如果聚簇轮廓模糊但准确率尚可通常是嵌入维度偏高导致低维投影失真这不算模型问题而是可视化手段的边界。从那以后我每次换数据集复现原型网络都强制走一遍完整流程先查每个类别的样本数量下限再固定随机种子跑 100 个 episode 验证能收敛最后才把评估轮数拉满到 1000。这套习惯帮我省掉了大量“实验结果忽高忽低”的无效对比时间。希望这篇拆解能让你复制项目时少走几步弯路。本文还有配套的精品资源点击获取