简介：基于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(dim=1) # (n_way, z_dim) # 查询点与所有原型的欧氏距离平方 z_query = z_query.contiguous().view(-1, z_dim) dists = torch.cdist(z_query, prototypes, p=2).pow(2) logits = -dists # 距离越近分数越高 return logits, prototypes第一行x.view(self.n_way, self.n_support, z_dim)成立的前提是输入数据的顺序完全符合采样器输出的顺序：先按类排列，再按“支撑在前、查询在后”排列。如果你自己实现数据加载器时打乱了样本顺序，这里会直接报维度错误，而且报错信息往往不容易一眼看出来。
torch.cdist是 PyTorch 里计算批量距离矩阵的函数，p=2表示欧氏距离。这里取了平方再取负作为 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(), lr=1e-3, weight_decay=1e-4) # cosine 退火：让学习率随训练进度平滑下降 scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=total_episodes, eta_min=1e-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 memory，batch 调小仍不行。
原因：原型网络虽然模型轻量，但嵌入向量会在每个 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_graph=True，并在训练循环里及时调用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_sampler：
import glob import os import torch from PIL import Image from torch.utils.data import Dataset class CustomEpisodicDataset(Dataset): def __init__(self, root, transform=None): # 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, padding=4), 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_size=128): model = models.resnet18(pretrained=True) # 去掉分类头，保留特征提取部分 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), dim=2 ) logits = sim / 0.1 # 除以温度系数，控制概率分布的锐度这里除以 0.1 是经验值，温度越高分布越平，梯度越平缓。如果换成余弦距离，原代码里的.pow(2)和取负都要删掉，否则符号彻底反了。马氏距离要引入协方差矩阵，小样本下每类只有 K 个样本根本估不准协方差，反而引入额外的偏差，除非你的每类样本量能达到 20 个以上，否则不建议尝试。
6. 让模型在真实场景里站得住：配方化评测的验证习惯
少样本模型能不能用，不能只靠一条准确率曲线判断，这是我踩过最大的坑。第一次训练时验证 loss 很漂亮，觉得自己调参技术不错，结果放到新类别上直接掉了一截。后来养成了一个习惯：固定随机种子，固定任务集，把评测变成“配方化”流程，每次改动都有可比基线。
具体做法是在验证脚本中加入三段固定逻辑。第一段固定种子，让每次试验的任务采样完全一致；第二段固定任务列表，在数据集里预先抽取一批评测任务，重复使用，保证不同模型配置之间的对比不受随机波动影响；第三段在报告里按 shot 分档，分别输出。下面是一个最小可复现的验证片段：
def evaluate(model, val_loader, n_way, k_shot, n_query, tasks=200): 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(dim=1) 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 两个档位的结果都记录下来，不再凭一条曲线判断模型好坏。希望帮到你。
本文还有配套的精品资源，点击获取
