简介基于PyTorch实现的原型网络少样本学习代码包面向机器学习初学者与少样本学习研究者完整覆盖原型网络论文中的核心方法与实验流程。资源共12个文件以7个Python脚本为主分别承担参数解析、原型网络模型构建、损失计算、批次采样、数据集加载及训练入口等功能并附带开源许可文件、说明文档和示例图像压缩包仅135KB结构紧凑便于对照源码进行学习。目前已有920人学习下载适合作为少样本学习方向的入门参考。代码在src目录中按模块拆分各脚本职责分明可在Omniglot数据集上直接运行分类实验说明文档与图像进一步展示了运行结果和依赖环境有助于理解原型网络的工作原理也方便在此基础上修改和扩展是兼顾可读性与实用性的轻量实现。1. 小样本学习与原型网络这份代码包到底能让你少踩多少坑做视觉任务的都懂一个残酷现实生产环境里能标注的图往往就几十张良品多不良品少新类一出现更是只有几张样张。传统分类模型在这个量级上基本是废的而小样本学习few-shot learning这类方法就是冲着“几张图也把类分出来”去的。原型网络Prototypical Networks是这条技术路线上最朴素也最稳的一支把每个类压缩成一个原型向量查询样本谁离得近就归谁。这份 PyTorch 代码包把论文中的训练逻辑完整落成了可跑的工程覆盖数据采样、模型搭建、损失计算、训练验证全套流程适合刚入手少样本学习的开发者也适合要先在自制小数据集上跑一条基线的同学。花两三天自己从零写不如直接基于这份已验证过的实现去改。2. 看懂原型网络的设计与代码结构类原型、欧氏距离与文件职责2.1 为什么“算均值”这件事在小样本下站得住原型网络的核心假设是存在一个嵌入空间同类样本在这个空间里聚成一团。那么每个类只需要一个代表点——原型prototype取支撑集support set样本嵌入的均值即可。查询样本query sample的类别就由它到每个原型的欧氏距离决定距离最近的类胜出。这个设计高明在两点。第一均值作为原型是一个低方差估计器即使每类只有 5 个样本算出来的类中心在嵌入空间里也相对稳定第二它把度量学习变成了“与类中心比较”不再需要像孪生网络那样两两配对比较计算量从 O(n²) 降到了 O(n)训练和推理都快很多。训练时用的损失不是简单的交叉熵而是把距离转换成概率再做负对数似然p(y k | x) exp(-d(f(x), c_k)) / Σ_j exp(-d(f(x), c_j))其中 d 是欧氏距离的平方c_k 是第 k 类的原型。这里有个容易误解的点距离越近概率越高所以 logits 取的是负距离。见过不少人在这一步把正负号搞反训练半天 loss 不降反升。2.2 代码包的典型结构与每个文件的职责这套实现的代码结构很清晰按职责拆分成了几个模块。拿到手先别急着跑花十分钟把文件对应关系看明白后面改起来才有方向├── prototypical_batch_sampler.py # 按 N-way K-shot 逻辑采样支撑集与查询集 ├── prototypical_networks.py # 编码器 原型分类头 ├── prototypical_loss.py # 原型损失负对数概率 ├── OmniglotDataset.py # Omniglot 数据加载与预处理 ├── utils.py # 工具函数模型保存、日志、训练流程 ├── train.py # 训练入口 └── val.py # 验证入口数据采样器是整个流程的起点。它保证了每个 batch 里同时包含支撑集和查询集。支撑集用来计算原型查询集用来计算损失和更新参数两者必须在类别层面严格对齐否则训练过程会退化成普通分类。2.3 核心代码读法采样器与网络前向采样器的逻辑关键是“先选类再选样本”。每次构建一个任务时从全部训练类中随机抽 N 个类N-way每个类里再抽 K 个支撑样本K-shot和 Q 个查询样本。代码结构大致如下# 伪代码每个 episode 的任务采样 def sample_episode(dataset, n_way, n_support, n_query): # 1. 从数据集类别中随机抽取 n_way 个类 selected_classes random.sample(dataset.classes, n_way) # 2. 每个类内分别抽支撑样本和查询样本 support [] query [] for cls in selected_classes: idxs dataset.class_to_indices[cls] n_support n_query min(n_support n_query, len(idxs)) chosen random.sample(idxs, n_support n_query) support.append(chosen[:n_support]) query.append(chosen[n_support:]) # 3. 打平后返回每类前 n_support 个是支撑后 n_query 个是查询 support [i for sub in support for i in sub] query [i for sub in query for i in sub] return support, query注意chosen[:n_support]和chosen[n_support:]这种切分方式它隐含了一个约定支撑集和查询集来自同一批标注数据但绝不能重叠。一旦查询样本在训练阶段被用于计算原型相当于模型在“开卷考试”验证指标会虚高到没有参考价值。网络前向部分则更为直接输入是拼接好的支撑与查询图像先过编码器得到嵌入再算原型和距离def forward(self, x): # x: (n_way * (n_support n_query), channels, h, w) z self.encoder(x) # 提取嵌入向量 z_dim z.size(-1) # 按支撑/查询切分嵌入 z_support z.view(self.n_way, self.n_support, z_dim) z_query z.view(self.n_way, self.n_query, z_dim) # 类原型 支撑集嵌入的均值 prototypes z_support.mean(dim1) # (n_way, z_dim) # 查询点与所有原型的欧氏距离平方 z_query z_query.contiguous().view(-1, z_dim) dists torch.cdist(z_query, prototypes, p2).pow(2) logits -dists # 距离越近分数越高 return logits, prototypes第一行x.view(self.n_way, self.n_support, z_dim)成立的前提是输入数据的顺序完全符合采样器输出的顺序先按类排列再按“支撑在前、查询在后”排列。如果你自己实现数据加载器时打乱了样本顺序这里会直接报维度错误而且报错信息往往不容易一眼看出来。torch.cdist是 PyTorch 里计算批量距离矩阵的函数p2表示欧氏距离。这里取了平方再取负作为 logits相当于对距离做了 softmax 归一化。从这行可以看出原型网络对距离度量极其敏感后面换自定义特征时这一行往往要反复调。3. 跑通训练与验证流程数据集、命令与关键参数的一次到位3.1 数据集准备Omniglot 与 emnist 的下载与路径处理这份代码默认支持两个数据集Omniglot 和 emnist。Omniglot 是手写字符集有 50 种字母体系、1623 类字符每类仅 20 个样本本身就是为小样本任务设计的emnist 是扩展版字母数字集类别更多但样本也更充足。数据准备阶段最容易出问题的地方是目录结构和路径配置。常见的做法是建立一个./data目录里面按omniglot/、emnist/分别存放。代码里读取数据的部分一般会要求目录下已经分好train/和val/子目录每个子目录下再按类别建文件夹。如果你下载的是官方原始压缩包通常需要先跑一次解压与整理脚本把松散的文件规整成上述结构。建议在首次运行前手动检查一下每个类文件夹下的图片数量。以 Omniglot 为例如果某个类不足 20 张而你在训练配置里把n_support n_query设到了 20 以上采样器会直接越界报错。我一般会在数据集加载时打印出类别数与最少样本数从源头把这种异常暴露出来。3.2 训练与验证命令从零到可复现基线环境准备好之后训练是几条命令就能跑起来的。以在 Omniglot 上训练 5-way 1-shot 为例# 训练 5-way 1-shot 模型每类包含 15 个查询样本 python train.py \ --n_way 5 \ --k_shot 1 \ --n_query 15 \ --episode 1000 \ --lr 1e-3 \ --data_root ./data \ --dataset omniglot \ --embedding_size 64这里每个参数都有明确的目的n_way 5表示每轮任务随机抽 5 个类k_shot 1表示每类只取 1 个支撑样本计算原型——这正是“少样本”的关键压力位n_query 15提供足够多的查询样本计算稳定的梯度episode 1000是训练任务总数而不是迭代轮数这点和普通分类训练差别很大。训练完成后再跑验证# 使用验证集评测加载最优模型权重 python val.py \ --model_path ./checkpoints/best_model.pt \ --n_way 5 \ --k_shot 1 \ --n_query 20 \ --episode 500 \ --dataset omniglot验证时的n_way和k_shot必须和训练保持一致否则比较没有意义。如果你想看模型在 5-shot 下的表现需要单独训练一个 5-shot 模型或者用训练好的 1-shot 模型做冷启动微调而不是直接拿 1-shot 模型去验证 5-shot 的准确率。3.3 优化器与调度器怎么搭更稳小样本训练和常规分类有个明显差异任务数量少、每任务内样本少梯度噪声大。因此优化器的选择对收敛质量影响显著。这份代码里常见的组合是 Adam 或 AdamW 配合余弦退火调度器下面是一个经过简单验证的配置import torch from torch import optim # AdamW 对嵌入空间训练更稳weight_decay 不宜过大 optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) # cosine 退火让学习率随训练进度平滑下降 scheduler optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxtotal_episodes, eta_min1e-5 )T_max要设成总的 episode 数而不是简单的 epoch 数。因为采样器每轮生成一个新任务批量天然较大学习率变化太快会导致模型还没在嵌入空间站稳就开始剧烈震荡。我自己常用的做法是先把lr固定在 1e-3 附近确认 loss 曲线有下降趋势后再引入余弦退火。训练过程中要同时关注train_loss与val_loss两条曲线。原型网络比较特殊的是训练任务是随机采样的所以训练 loss 本身波动就不小不要因为几个 episode 的跳动就慌张调整参数。看趋势至少要看五十到一百个 episode 的滑动平均。4. 常见问题与排查五个会把模型训练带偏的坑4.1 数据泄漏训练时 loss 低得离谱验证时原形毕露现象训练 loss 一路跌到 0.1 以下验证准确率却只有 50% 上下两者差距大到不像是同一种模型。原因支撑集和查询集采样时发生了重叠。常见触发点有两个——一是数据加载器的 shuffle 逻辑把同类的支撑和查询样本混在一起了二是自定义数据集时同一张图片既被选中到支撑集又被选中到查询集。解决检查采样器代码确保random.sample一次性抽完n_support n_query个索引后再切分而不是分两次独立抽样。另外在训练脚本里加一个断言对每个任务的支撑集和查询集索引求交集非空就终止# dtype 检查支撑集与查询集不应有交集 assert len(set(support_idx) set(query_idx)) 0, \ Support set overlaps with query set!4.2 验证配置不一致训练 1-shot 却拿 5-shot 的成绩自嗨现象验证准确率有 85%换到真实场景每类只有 1 张图后只有 60%性能掉了一截。原因验证脚本里的--k_shot设置和训练配置不一致。1-shot 训练的模型对单样本原型的估计已经形成习惯验证时突然给 5 个样本算均值原型相当于降低了任务难度准确率自然会好看。解决默认规则是“训练用什么配置验证就原样用什么配置”。如果要发布三组结果如 1-shot、5-shot、10-shot每一组都必须单独训练或基于对应配置评估。把验证配置做成从训练日志里读取而不是手工重新传参能从机制上杜绝这类不一致。4.3 显存溢出图片分辨率与 embedding_size 的双重夹击现象训练刚开始就报CUDA out of memorybatch 调小仍不行。原因原型网络虽然模型轻量但嵌入向量会在每个 episode 被整体复制。如果输入图片分辨率大比如 224×224加上embedding_size设到 256 以上一个 batch 里的支撑和查询样本总数很容易把显存吃满。解决先降分辨率。原始 Omniglot 图像是 28×28这套代码默认也是这个尺寸迁移到自己的数据时最省显存的做法是把所有图统一缩放到 84×84 或 64×64。再不行就把n_query从 15 降到 5因为支撑集大小是固定死的查询样本数量才是可控项。显存消耗公式大约是(n_way * (k_shot n_query)) * embedding_size按这个公式预先算好预算再调参。4.4 损失模块与我们常见的 PyTorch 版本不兼容现象运行到loss.backward()时抛出一个AttributeError提示某个损失内部属性不存在例如class_loss或retain_graph相关参数不被识别。原因这份代码大约编写于较早期的 PyTorch 生态训练时经常会把两个损失实例训练损失和验证损失放在同一个 forward 里交替调用。新版 PyTorch 对损失计算图的保留语义做了调整旧的写法在一些版本下会崩。解决最简单的方法是把train_loss和val_loss两个损失实例合并成同一个每次前向调用后立即清空计算图。如果坚持保留两个实例就在前向时显式传入retain_graphTrue并在训练循环里及时调用torch.cuda.empty_cache()释放中间缓存。4.5 学习率调度器的“玄学”loss 卡在一个平台上下不来现象训练两三万个 episode 后loss 稳定在 0.6 附近怎么调都下不去。原因调度器的T_max设错了。如果T_max设成 epoch 数而不是 episode 数而一个 epoch 里包含几百个 episode那么余弦退火的周期和实际训练节奏完全错位学习率要么衰减太快要么迟迟不衰减。解决把T_max明确设置为训练总 episode 数并且在日志里每个调度步打印一次当前学习率。如果你已经卡在平台上可以直接把学习率强制降一个数量级如从 1e-3 降到 1e-4再接着训练很多时候就能继续往下走。这类问题没有通用解但“降低学习率重跑一段”是成本最低的后悔药。5. 把模型用到自己的数据上从 Omniglot 迁移到自定义图像集5.1 自定义数据集接入的三个改动点自己的数据往往不是 28×28 的灰度图而是几百张 RGB 图片。接入这套代码最核心的是替换数据加载部分保持输出格式仍是“每类一个目录、每张图一个文件”同时完成三步改动。第一步写一个继承自Dataset的类把每张图的类别索引直接传给prototypical_batch_samplerimport glob import os import torch from PIL import Image from torch.utils.data import Dataset class CustomEpisodicDataset(Dataset): def __init__(self, root, transformNone): # root 下每个子目录代表一个类 self.classes sorted(os.listdir(root)) self.class_to_indices {} self.paths [] self.labels [] for cls in self.classes: imgs glob.glob(os.path.join(root, cls, *.jpg)) idx len(self.paths) for p in imgs: self.paths.append(p) self.labels.append(cls) self.class_to_indices[cls] list(range(idx, len(self.paths))) self.transform transform def __len__(self): return len(self.paths) def __getitem__(self, idx): img Image.open(self.paths[idx]).convert(RGB) if self.transform: img self.transform(img) return img, self.labels[idx]这里class_to_indices是关键它让采样器能按类快速取索引而不是每次全表扫描。__getitem__返回的 label 在原型网络里其实不会被当作监督信号使用真正的监督来自支撑集原型的自监督式对比所以 label 类型是字符串还是整数都无所谓但必须和采样器里的类名一致。第二步把 transform 换成适配自己数据尺寸的组合。常见做法是把图片缩放到 84×84再做随机裁剪和归一化这一步是嵌入式空间质量的直接影响因素transform transforms.Compose([ transforms.Resize((84, 84)), transforms.RandomCrop(84, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])这里的均值方差直接沿用了常见预训练模型的统计值不要改成自己数据集的统计量。否则编码器初始化的语义空间会被破坏小样本场景下没有足够样本去重新校准。第三步在训练脚本里把数据集的加载入口换成CustomEpisodicDataset。如果你的图片分辨率超过了 84×84注意第 4 节说的显存问题先把图上采样成本降下来。5.2 backbone 怎么换从几层卷积到残差网络原本的编码器一般是四层卷积堆叠对 28×28 的手写字符足够但对自然图像明显太浅。换上预训练的残差网络是最直接有效的一步但要注意最后一层输出的 embedding_size 必须与损失函数里的距离计算维度保持一致。常见做法是把预训练特征提取器的全连接层替换成一层 1×1 卷积或线性映射输出维度控制在 64~128 之间。一个成熟的替换方案import torchvision.models as models def build_encoder(embedding_size128): model models.resnet18(pretrainedTrue) # 去掉分类头保留特征提取部分 model.fc torch.nn.Linear(512, embedding_size) return model换成残差网络之后embedding_size 从 64 提升到 128模型容量大了但训练需要的 episode 数量也要相应增加。以我自己的迁移经验原来 1000 个 episode 能收敛加大 backbone 后至少需要 2000 个 episode 起步。另外注意预训练模型的权重在编码器部分已经很成熟了训练时要么冻结前面几层要么给 backbone 一个较小的学习率如 backbone 1e-4、新加层 1e-3否则刚起步几个 episode 就会把这个预训练空间打乱。5.3 距离度量想换成别的余弦距离与马氏距离的改法欧氏距离在原型网络里是默认选项但它有个隐含假设嵌入空间各方向等权。当你把 backbone 换成了残差网络嵌入向量的尺度范围会因预训练权重而变化此时余弦距离往往更稳。改动点很小直接用torch.nn.functional.cosine_similarity替换距离计算同时把前向里的 logits 从“负距离平方”改成“余弦相似度”# 将 z_query 与 prototypes 做余弦相似度计算 sim torch.nn.functional.cosine_similarity( z_query.unsqueeze(1), prototypes.unsqueeze(0), dim2 ) logits sim / 0.1 # 除以温度系数控制概率分布的锐度这里除以 0.1 是经验值温度越高分布越平梯度越平缓。如果换成余弦距离原代码里的.pow(2)和取负都要删掉否则符号彻底反了。马氏距离要引入协方差矩阵小样本下每类只有 K 个样本根本估不准协方差反而引入额外的偏差除非你的每类样本量能达到 20 个以上否则不建议尝试。6. 让模型在真实场景里站得住配方化评测的验证习惯少样本模型能不能用不能只靠一条准确率曲线判断这是我踩过最大的坑。第一次训练时验证 loss 很漂亮觉得自己调参技术不错结果放到新类别上直接掉了一截。后来养成了一个习惯固定随机种子固定任务集把评测变成“配方化”流程每次改动都有可比基线。具体做法是在验证脚本中加入三段固定逻辑。第一段固定种子让每次试验的任务采样完全一致第二段固定任务列表在数据集里预先抽取一批评测任务重复使用保证不同模型配置之间的对比不受随机波动影响第三段在报告里按 shot 分档分别输出。下面是一个最小可复现的验证片段def evaluate(model, val_loader, n_way, k_shot, n_query, tasks200): torch.manual_seed(42) np.random.seed(42) random.seed(42) accs [] for _ in range(tasks): support, query sample_episode( val_loader.dataset, n_way, k_shot, n_query ) logits, prototypes model(val_loader.dataset[support query]) pred logits.argmax(dim1) acc (pred torch.arange(n_way).repeat_interleave(n_query)).float().mean() accs.append(acc.item()) return { mean_acc: float(np.mean(accs)), std_acc: float(np.std(accs)), }注意这里sample_episode用的就是训练时的同一个采样函数进一步保证了训练和验证的数据分布一致。保存模型时也要按这个思路做best_model.pt保存验证集上 mean_acc 最优的权重而不是最后一个 episode 的权重。对少样本任务来说最后阶段的模型参数往往已经在超调边缘早停在训练小样本模型时是默认操作而不是备选项。最后一层防护是给模型做“单样本压力测试”也就是在真实部署前测试模型在只有 1 张图支撑时的表现。这是我在一次安防类项目里吃了大亏后才彻底养成的习惯——训练时 n_query 多、梯度稳部署时每类只有一张图loss 和准确率立刻变脸。从那以后我每次训练少样本模型都强制跑一遍配方化评测顺便把 1-shot、5-shot 两个档位的结果都记录下来不再凭一条曲线判断模型好坏。希望帮到你。本文还有配套的精品资源点击获取