简介针对小样本学习场景这份基于PyTorch的开源项目完整实现了Prototypical Networks原型网络源自arXiv 1703.05175论文代码覆盖从数据加载、批次采样、模型推理到损失更新的全流程适合正在入门元学习、少样本分类的研究者或工程师参考。包内共12个文件其中7个Python脚本构成核心代码分别负责模型结构定义、原型损失计算、批采样策略、Omniglot数据集加载、命令行参数解析与训练主流程另含2张示意图、README说明与LICENSE许可压缩包仅135KB体积小巧、目录层级清晰便于逐模块阅读和二次开发。目前已有920人学习下载。借助该项目可以快速在Omniglot上跑通few-shot实验理解类原型表征与度量学习的实现细节也能以train.py为主入口调整参数、替换数据或扩展损失函数搭建自己的小样本学习基线方案适合作为论文复现和课程设计的起点。1. 小样本分类的经典解法Prototypical Networks 用“类代表”完成 5-shot 识别假设你现在手里只有一个新类别的 5 张照片传统分类网络会当场失效softmax 输出层是针对训练时见过的类别定死的新类进来既没有对应神经元也没有足够数据重训。Prototypical Networks原型网络就是 few-shot learning 里绕不开的经典方案它把每个类别压缩成一个「原型向量」新的输入图像在嵌入空间里跟所有原型比距离离谁近就判给谁整个推理过程不需要为新类做任何梯度更新。这个 PyTorch 项目正是该方法的完整实现从 episode 采样、卷积 backbone、原型计算到训练评估一应俱全拿到手就能在 Omniglot、miniImageNet 这类小样本基准上复现 5-way 1-shot 和 5-way 5-shot 的实验结果。适合两类人正在做小样本图像分类课题的学生以及刚入手 few-shot 想找一个干净参照实现的 Python 开发。2. 项目全景与技术选型一个 episode 从数据采样到 loss 的完整链路打开这类仓库之前建议先理解 few-shot 训练跟普通图像分类的本质区别。普通分类每个 batch 从所有类别里均匀采样图片模型最终输出一个长度为「总类别数」的向量few-shot 则把每一次训练变成一个「小任务」每个小任务里只有 N 个类别、每类 K 张支撑图片模型要在这个临时任务上学会区分。这个「小任务」就是 episode整个项目的一切代码都在围绕它转。2.1 仓库结构每个文件在 few-shot 流程里的位置我拆到手的仓库不大核心代码集中在 src 目录入口脚本放在 scripts 下典型的组织方式如下. ├── data/ │ ├── omniglot/ # 手写字符数据集原始文件 │ └── mini_imagenet/ # miniImageNet 按类目录整理 ├── src/ │ ├── model.py # PrototypicalNetwork 与 backbone │ ├── dataset.py # 数据集读取、episode 采样器 │ ├── trainer.py # 单 episode 训练逻辑 │ └── utils.py # 参数解析、seed 固定、日志 ├── scripts/ │ ├── train.py # 训练入口 │ └── evaluate.py # 评估入口 ├── checkpoints/ # 训练产出的模型权重 └── README.md各文件的职责和阅读顺序可以用一张表说清楚文件职责建议阅读顺序src/model.py卷积嵌入网络、原型计算、距离打分1src/dataset.py图片加载、support/query 划分、EpisodeSampler2src/trainer.py损失计算、反向传播、指标统计3scripts/train.py命令行参数、数据加载、训练主循环4scripts/evaluate.py恢复 checkpoint、固定 seed、跑 N 个 episode5先从 model.py 读起是对的因为原型网络的核心只有「嵌入 距离」两个概念模型文件通常不超过 100 行dataset.py 反而是最容易出 bug 的地方support 和 query 的索引一旦错位训练曲线会非常难看。2.2 技术选型拆解PyTorch、episodic 训练和欧氏距离这个项目选 PyTorch 是有原因的。episode 采样本质上要求「每一轮重新组合一批图片」数据集类的__getitem__返回的应该是一个索引而不是固定图片这种动态图风格在 PyTorch 里写起来最顺手。调试时你可以在训练中途随手打印 logits 的张量形状和数值范围不需要像静态图时代那样把整个计算图提前编译好。backbone 通常选 4 层卷积每层 64 个 3×3 卷积核配 BatchNorm、ReLU 和 2×2 最大池化。小样本任务的训练数据本来就少几十层的 ResNet 在这类任务上很容易过拟合四层 CNN 参数少、收敛快在 Omniglot 这种 28×28 的小图上已经能压到论文同档水平。输入通道数通过参数控制跑 Omniglot 传1跑 miniImageNet 传3。距离度量默认是欧氏距离的平方这也不是随手选的。原型网络有一个理论保证当嵌入空间的距离度量是欧氏距离时类内样本的均值向量就是该类别的最优原型表示。一旦你换成余弦距离或者马氏距离「类均值 最优原型」这个性质就不成立了效果需要重新调。这也是为什么很多教程劝你别一上来就改距离度量我会在最后一章专门讲怎么改才稳。2.3 episode 采样逻辑support 和 query 是怎么切出来的dataset.py 里最核心的组件是 EpisodeSampler它的作用就是每轮从训练集中动态抽出一个任务。我见过的大部分实现核心逻辑都长这样class EpisodeSampler: def __init__(self, labels, n_way, k_shot, query_shot): self.labels torch.tensor(labels) self.class_ids torch.unique(self.labels) self.n_way n_way self.k_shot k_shot self.query_shot query_shot def sample(self): chosen self.class_ids[torch.randperm(len(self.class_ids))[:self.n_way]] support_ids, query_ids [], [] for c in chosen: idx torch.where(self.labels c)[0] perm idx[torch.randperm(len(idx))] support_ids.append(perm[:self.k_shot]) query_ids.append(perm[self.k_shot:self.k_shot self.query_shot]) return torch.cat(support_ids), torch.cat(query_ids)这段代码做了三件事先从全部类别里随机选n_way个类再在每个类内部随机打乱样本索引最后把前k_shot个分给 support、接下来query_shot个分给 query。注意 support 和 query 来自同一个类别集合但不重叠这样才能保证「用 K 张样本学到区分」这个任务成立。实际使用时有一个硬约束数据集中每个类别至少要有k_shot query_shot张图片否则索引会越界。Omniglot 每类至少有 20 张样本所以 5-shot 配 15 个 query 没问题你自己收集的数据集如果每类样本不足 20 张就得把 query_shot 相应调小或者做数据增强扩量。3. 核心模块拆解从 image tensor 到原型向量再到 logits这一章是整个仓库的发动机。原型网络的前向过程可以拆成三步嵌入函数把图片映射到特征空间每个类的 support 特征求均值得到原型query 特征和所有原型算负平方欧氏距离得到 logits。走完这三步剩下的就是交叉熵损失和梯度更新。3.1 嵌入网络四层卷积块搭起来的特征提取器model.py 里先定义了一个可复用的卷积块然后把它堆四次形成嵌入网络import torch.nn as nn class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) def forward(self, x): return self.net(x) class EmbeddingNet(nn.Module): def __init__(self, in_ch1, hidden64): super().__init__() self.encoder nn.Sequential( ConvBlock(in_ch, hidden), ConvBlock(hidden, hidden), ConvBlock(hidden, hidden), ConvBlock(hidden, hidden), ) def forward(self, x): return self.encoder(x)这里每个 ConvBlock 都是「卷积 BN ReLU 池化」四件套。卷积用 3×3、padding1 保证尺寸不缩真正缩小分辨率的是 MaxPool2d(2)每层把宽高减半。BN 的作用在小样本场景里比较微妙它对每个 batch 统计均值方差episode 内样本数量少时统计量噪声大后面避坑章我会专门讲这个问题。参数说明in_ch是输入图像的通道数Omniglot 灰度图传 1miniImageNet 彩色图传 3hidden是每层卷积的输出通道数常用 64显存紧张可以降到 32效果会略掉。28×28 的输入经过四层池化后会变成 1×1 特征图这其实是好事后续处理非常清爽84×84 的输入会得到 5×5 的特征图处理方式稍有不同。3.2 原型计算与距离打分类均值加负平方欧氏距离模型的主干逻辑在 PrototypicalNetwork 里support 和 query 共用同一个嵌入网络然后分别处理import torch import torch.nn.functional as F class PrototypicalNetwork(nn.Module): def __init__(self, in_ch1, hidden64): super().__init__() self.embedding EmbeddingNet(in_ch, hidden) def forward(self, support, query, n_way, k_shot): support_feat self.embedding(support) query_feat self.embedding(query) # 原型 每类的 support 特征求均值 support_feat support_feat.view(n_way, k_shot, -1) prototypes support_feat.mean(dim1) # query 特征展平和原型做欧氏距离 query_feat query_feat.view(query.size(0), -1) dists torch.cdist(query_feat, prototypes, p2) ** 2 return -distsview(n_way, k_shot, -1)这一步的含义是把每个类的k_shot个样本特征叠在一起最后的-1是嵌入维度在 28×28 输入下就是 64。mean(dim1)对第 1 维求均值得到(n_way, 64)形状的原型矩阵。对于 84×84 输入特征图是(B, 64, 5, 5)view(-1)会直接展平成 1600 维相当于把空间位置也当成特征更省参数的做法是先接一个自适应平均池化压到 64 维两种写法在这个仓库里都见过。torch.cdist(query_feat, prototypes, p2)计算的是欧氏距离再平方是因为论文里的距离项是-||x - c||^2。合起来代码返回的 logits 形状是(query_total, n_way)每一行表示这个 query 样本和 N 个类原型的相似度打分。距离越小负值越小对应的类概率越大。3.3 训练逻辑负对数似然当作 loss一次 episode 一次更新原型网络没有单独的分类头也不需要在 episode 内做任何梯度更新所以训练循环非常短import torch.nn.functional as F def train_one_episode(net, optimizer, support, query, query_labels, n_way, k_shot): net.train() logits net(support, query, n_way, k_shot) loss F.cross_entropy(logits, query_labels) acc (logits.argmax(dim1) query_labels).float().mean().item() optimizer.zero_grad() loss.backward() optimizer.step() return loss.item(), acc这里的query_labels是一个长度为query_total的索引数组取值范围是0到n_way - 1表示每个 query 样本在本次 episode 中属于第几个类。F.cross_entropy内部自带 softmax所以前面返回-dists直接传进来即可不需要再手动做归一化。值得注意的细节是梯度流向。loss 对 logits 求梯度logits 又由 query_feat 和 prototypes 共同决定而 prototypes 来自 support_feat 的均值因此梯度会同时流过 query 分支和 support 分支。也就是说support 图像虽然没有直接参与损失计算但它们的原型参与了梯度反传嵌入网络照样能从 support 侧学到表征。训练脚本里还需要在每轮 episode 结束后打印 loss 和 acc。刚开始训练时 acc 应该在 1/n_way 附近比如 5-way 任务随机猜就是 20%随着训练进行缓慢爬升。如果 loss 一上来就卡在某个奇怪数值上不动优先检查 logits 的量级正常应该在 0 附近取负绝对值不超过十几。4. 训练与评估在 Omniglot 和 miniImageNet 上跑出实验数字模型代码看懂了接下来就是把训练跑起来。这个仓库的入口都在 scripts 目录下数据准备、训练参数、评估方式三条线分开互不干扰。第一次跑强烈建议先用 Omniglot 验证环境成功后再上 miniImageNet。4.1 数据准备两类数据集的预处理差异Omniglot 是一个手写字符数据集每张图是灰度图原始尺寸 105×84常见做法是缩放到 28×28 并做四方向旋转增强0、90、180、270 度相当于把数据量扩大四倍。数据划分要严格使用官方自带的 background 和 evaluation 分割训练类、验证类、测试类互不重叠否则评估结果没有说服力。miniImageNet 是 ImageNet 的一个子集100 个类别共 6 万张图片常用划分是 64 类训练、16 类验证、20 类测试。图像是彩色图常见做法是缩放到 84×84训练时加随机裁剪和水平翻转增强验证时不加。预处理上的差异对结果影响非常明显很多复现数字对不上都是卡在归一化方式和 crop 细节上。准备好数据后训练命令长这样python scripts/train.py \ --dataset omniglot \ --n-way 5 \ --k-shot 5 \ --query-shot 15 \ --train-episodes 30000 \ --val-episodes 200 \ --optimizer adam \ --lr 0.001参数说明--n-way 5表示每个 episode 里有 5 个类别这是 few-shot 里最常用的配置--k-shot 5表示每类给 5 张支撑图片如果想做 1-shot 实验改成 1 即可--query-shot 15表示每类还要额外抽出 15 张查询图片用来计算损失。--train-episodes就是训练总轮数小样本任务的「epoch」概念很弱一般直接用 episode 数来控制训练量。4.2 训练参数配置不同数据集下的常用组合我把两类数据集上常见的参数组合整理成一张表方便对照参数Omniglot 常用值miniImageNet 常用值n_way55k_shot1 / 51 / 5query_shot1515输入尺寸28×28 灰度84×84 彩色训练 episode 数10000 3000020000 60000优化器Adamlr 1e-3SGD momentum 0.9lr 0.1 衰减批量 episode 数32 或 6432 或 64Omniglot 任务相对简单Adam 配 1e-3 足够快到收敛miniImageNet 更难很多复现会用 SGD 配上学习率衰减效果比 Adam 更稳。注意「批量 episode 数」跟前面的n_way无关它表示一次梯度更新同时跑多少个 episode通常取 32太大会让每个 step 变慢太小则 BN 统计量不稳定。训练过程里我习惯每 1000 个 episode 打印一次 loss 和 acc同时把模型 checkpoint 存到 checkpoints 目录。第一次跑建议开一个 10000 episode 的小规模实验确认曲线正常后再放长时间不然发现配置错了前面几小时全白跑。4.3 评估脚本准确率之外一定要看置信区间评估阶段和训练阶段最大的区别是评估要固定 seed并且跑很多个 episode 取平均。few-shot 的单次 episode 结果方差非常大可能这次抽到的任务全是容易区分的类准确率 90% 以上下次抽到长得相似的大类准确率直接掉到 50%。因此评估命令里通常有一个--eval-episodes参数python scripts/evaluate.py \ --checkpoint checkpoints/omniglot_5way5shot.tar \ --dataset omniglot \ --n-way 5 \ --k-shot 5 \ --eval-episodes 1000 \ --seed 2024评估脚本的核心逻辑是遍历指定数量的 episode累加每次预测正确的样本数最后除以总查询样本数得到平均准确率。数据集图片少的时候可以每类重复多抽几个 query保证分母足够大。跑完 1000 个 episode 后我一般会顺手算一下 95% 置信区间import math def confidence_interval(acc, n_episodes): return 1.96 * math.sqrt(acc * (1.0 - acc) / n_episodes)n_episodes越大区间越窄。100 个 episode 下区间可能在 ±5% 到 ±8%1000 个 episode 能把区间压到 ±1% 左右。评测报告里如果只写一个准确率不给区间基本等于没写因为你无法判断两个模型之间的差异到底是真实差距还是采样波动。评估时不要用训练时的随机状态固定 seed 到 2024 这类值结果才可复现。5. 避坑指南训练不收敛与评估方差大的常见问题排查这里集中写几条我在复现这类仓库时实际踩过的坑全部按照「现象 → 原因 → 解决」拆开。前面 2.3 节的view和索引逻辑还有 3.3 节的标签对齐是最容易出错的两个地方下面都有对应记录。5.1 训练期loss 不降卡在某个固定数值附近现象训练刚开始loss 稳定在 1.6 附近5-way 任务的随机交叉熵就是 ln 5 ≈ 1.609跑了很久也不下降或者偶尔下降一点又弹回去。原因先判断是不是标签错位。query_labels 如果用了数据集的原始类别编号比如 0 到 99而 logits 只有 5 列F.cross_entropy不会报错但梯度方向完全错误loss 会一直停在一个诡异值。其次是学习率问题Adam 配 1e-3 跑 Omniglot 一般没问题但如果你换了更大的 backbone 或更深的网络梯度量级变大1e-3 可能导致震荡。解决先打印 logits 的形状和数值范围确认是(query_total, n_way)且数值在 -10 到 0 之间然后检查query_labels是不是 0 到 n_way - 1 的索引。两者都没问题就把学习率降到 1e-4 试跑 1000 个 episode曲线仍不降再回头看数据预处理。5.2 训练期精度突然掉到谷底恢复得很慢现象训练进行到中段acc 从 0.95 突然掉到 0.2之后要花很长时间才能爬回来甚至直接崩掉。原因最常见的是 BatchNorm 在小 episode 上统计量不稳定。训练进入中后期单 batch 的均值方差抖动变大BN 的 running stats 被带偏紧接着整个网络输出漂移。另外学习率过大时某个 batch 的梯度会一步把参数推出正常区域。解决如果用的是 32 个 episode 一起算梯度先把 batch 提到 64 增大 BN 统计样本数或者干脆把 BN 换成 GroupNorm小样本任务中 GroupNorm 往往比 BN 稳很多代价是收敛速度略慢。模型权重如果保存过中段 checkpoint可以回滚到掉点之前的版本调小学习率再继续。5.3 评估期验证准确率正常测试集准确率却惨不忍睹现象验证时 acc 稳定在 98% 左右换到测试类以后只剩 60%看起来模型完全不泛化。原因类划分泄漏这是 few-shot 实验里最隐蔽的翻车点。如果数据划分时随机把所有类混在一起分训练和测试同一个字符集的训练类可能来自同一批手写风格测试类跟训练类过于相似验证阶段自然准确真正面对没见过的字符集时风格差异一上来立刻打回原形。解决用数据集的官方背景划分。Omniglot 直接用 background 和 evaluation 两批数据背景字符集用来训练评估字符集用来测试miniImageNet 用标准的 64/16/20 划分。切数据之前先算一遍训练类和测试类的交集确保为 0这个检查我每次都会做。5.4 评估期置信区间 ±8%完全没法比较模型好坏现象两个模型跑评估一个 acc 0.62另一个 0.64但置信区间算出 ±0.08无法得出结论。原因评估的 episode 数太少。few-shot 单个任务方差本来就大如果只跑 50 个 episode采样波动完全盖过了模型差异。这算是评估设计的问题不是模型问题。解决把--eval-episodes提到至少 1000。如果数据集本身的类不够多可以降低 query_shot 并增加重复采样的轮数。评估时固定 seed让两个模型面对完全相同的任务序列这样比较才有意义。另外同时记录每个 episode 的 acc 分布而不只是平均能更早发现问题。5.5 玄学同一份代码别人 0.99自己只有 0.5现象跟着教程复现同样的配置别人的屏幕上是漂亮的收敛曲线自己这边 acc 一直在 50% 左右晃动代码看起来一模一样。原因这种「玄学」问题九成出在预处理上。比如 Omniglot 是否做了旋转增强像素归一化是除以 255 还是用 0.5 均值图片 resize 是 28×28 还是 32×32这些细节不对嵌入空间整体漂移距离度量全乱。解决对比实现时不要只看模型代码把数据 transform 逐行对一遍。打印输入 tensor 的 shape、dtype、数值范围很多代码块里默认输入是 float 0-1 区间如果你的输入是 uint8BN 和卷积的统计量会完全不对。从那以后我每次复现小样本项目第一件事永远是先用一小批数据跑通 forward打印中间张量确认数值量级合理再上完整训练这个习惯能省掉大量调参时间希望帮到你。6. 二次开发替换 backbone、距离度量与数据集的三板斧仓库跑通只是第一步大部分开发者拿这份代码是为了改造到自己的场景。三个最常动的点换特征提取网络、换距离度量、换自定义数据集。前两个都属于模型层改动第三个属于数据层改动。6.1 把四层 CNN 换成 ResNet 嵌入小模型跑通后想提升表征能力常见的做法是把 EmbeddingNet 换成预训练 ResNetimport torch.nn as nn import torchvision.models as models class ResNetPrototype(nn.Module): def __init__(self, backboneresnet18, embedding_dim128): super().__init__() base models.resnet18(weightsNone) base.conv1 nn.Conv2d(3, 64, kernel_size7, stride2, padding3, biasFalse) self.encoder nn.Sequential(*list(base.children())[:-1]) self.head nn.Linear(512, embedding_dim) def forward(self, x): feat self.encoder(x).flatten(1) return self.head(feat)注意两个细节第一替换 backbone 后输入的图像尺寸最好提到 224×224否则 ResNet 前几层的下采样会把小尺寸特征压没第二预训练权重对同分布数据有帮助但如果你的数据分布和预训练域差太远加载预训练权重反而会拖慢收敛。6.2 把欧氏距离换成余弦相似度有些任务下特征向量的长度包含无用信息用余弦相似度更合理。替换方式很直接import torch.nn.functional as F def cosine_logits(query_feat, prototypes, temperature10.0): q F.normalize(query_feat, dim1) p F.normalize(prototypes, dim1) return torch.mm(q, p.t()) * temperaturenormalize把向量长度归一为 1点积等价于余弦相似度。后面的temperature很重要归一化后点积最大值不超过 1直接做 logits 会让 softmax 太平、梯度消失乘以 10 把数值范围拉回来。换度量后必须重新调参建议固定 backbone、seed 和 episode 数先跑 5-way 1-shot 对比欧氏和余弦的收敛曲线。验证改造是否有效最直接的方法是同一份评估脚本分别跑改造前后的 checkpoint比较 1000 个 episode 的均值与置信区间。如果新度量掉点先检查 temperature 在 5 到 20 之间扫一遍再看有没有必要对特征做 L2 归一化前的缩放。我自己踩过很多次 cosine 掉点最后都是 temperature 没调对。希望这些改造思路能让你在这个仓库上少走几步弯路。本文还有配套的精品资源点击获取