PyTorch原型网络实战：少样本分类的轻量级方案与代码详解-北京尧图网络

简介这是一个基于PyTorch实现的Prototypical Networks原型网络少样本学习项目对应arXiv论文1703.05175主要面向机器学习初学者、算法研究人员以及需要在少量标注样本条件下完成图像分类任务的开发者。项目以Omniglot数据集为实验载体通过原型网络在嵌入空间中计算类原型并完成距离度量分类能够帮助读者快速理解少样本学习中的核心思路。资源为ZIP压缩包大小仅135KB共12个文件包括7个Python源码、2张PNG图片、1个Markdown说明文档、1个许可证文件和1个.gitignore配置Python源码覆盖模型定义、训练流程、损失函数、批次采样器和数据集处理等关键模块PNG图片提供网络结构或实验效果示意Markdown文档对项目使用方式进行说明。包体精简但结构清晰便于直接阅读和二次开发。目前已有920人学习下载适合作为PyTorch少样本学习入门的参考实现。1. Prototypical Networks五张图也能训出一个能用的分类器做机器学习落地的人大概率碰过这类需求新品类只有五到十张样例标注预算又只够标这么多却要模型立刻能认出它。Prototypical Networks原型网络是 Few-shot Learning 里代价最小的一条路——它不需要微调整个网络只要在嵌入空间里找到每个类的“原型”原型等于支持集样本的均值向量然后让查询样本去距离最近的原型认亲。这份用 PyTorch 实现的项目把训练、验证、推理拆得干干净净适合两类人一类是刚入门的 Python 开发想在小样本分类上跑通一个能复现的基线另一类是已经在做迁移学习的工程师想给现有特征提取器接一个轻量分类头。2. 技术原理与架构原型、距离与 Episode 三者怎么配合2.1 从“找代表”开始支持集、原型与查询集原型网络的思想非常朴素每个类别在嵌入空间里找一个“代表向量”这个代表向量叫做原型prototype。给定一个嵌入函数 f_phi通常是一个卷积网络第 k 类的原型就是该类所有支持集样本嵌入向量的均值c_k (1 / |S_k|) * Σ f_phi(x_i)其中 (x_i, y_i) ∈ S_k这个均值不是随便选的。在欧氏空间里类内样本的均值恰好是“最小化类内点到该点距离平方和”的解也就是说均值是该类在欧氏度量下的最优代表。这就是为什么原型网络在数学上成立——如果距离度量换了均值的“最优性”也就不成立了。查询样本 x 属于第 k 类的概率用负欧氏距离过 softmax 得到。距离越小、负值越大概率越高。整个模型没有额外的分类层分类器就是“查询点与各原型的距离比较器”所以训练好的网络可以随时接纳新类别只要你能算出新类别的原型。2.2 Episode 训练机制每个 batch 就是一次“小任务”传统分类训练里一个 batch 是从所有类别里均匀采样的图片模型一眼就能看到所有类别。而 Few-shot 场景下训练和测试必须保持一致——模型在训练时看到的也是“一个小任务”。这个项目采用标准的 Episode 训练机制每个 episode 从训练集的类别池里随机挑 N 个类别N-way每个类别取 K 张作为支持集K-shot再取 Q 张作为查询集N-query。模型只在这个小任务上计算损失、更新梯度下一个 episode 又换一批随机类别。这样反复训练嵌入网络学会的不是“记住某个类”而是“如何在任意一个新任务里把同类拉近、异类推开”。这种机制与对比学习Contrastive Learning有相似之处但更直接。每个 episode 天然就是一个 N 分类任务损失直接用交叉熵不需要构造正负样本对代码上少了很多弯弯绕绕。第一次接触的人不用怕它的实现比 triplet loss 还要简单核心代码加起来不到一百行。2.3 项目目录结构与训练流水线拿到这份 PyTorch 工程先看目录再动手。它的分层很清爽模型、数据、入口分开方便替换。prototypical_networks/ ├── configs/ │ └── default.yaml # 训练超参n_way / n_shot / n_query / lr ├── data/ │ └── episode_loader.py # Episode 采样器与数据集包装 ├── models/ │ ├── conv4.py # 嵌入网络 │ └── proto_net.py # 原型网络封装 ├── train.py # 训练入口 ├── evaluate.py # 验证入口 └── utils/ └── metrics.py # 准确率与损失统计训练流水线也清晰先从 dataset 里解析出“类别 → 样本索引”的映射然后每个 epoch 里反复调用 episode 采样器生成小任务把支持集和查询集都过一遍嵌入网络算原型、算距离、算损失反向传播更新网络。下面是这份工程最核心的超参表。参数含义常见取值影响n_way每个 episode 的类别数5 或 20越大任务越难需要更深网络n_shot每类支持集样本数1 或 51-shot 是难度上限n_query每类查询集样本数15 或 16影响梯度估计稳定性meta_batch_size每次迭代的 episode 数1显存有限时固定为 1lr学习率1e-3 到 1e-4用小网络时 1e-3 起步embedding_dim嵌入向量维度64维度太低表达能力受限我一般会先用 n_way5、n_shot1、n_query15 跑通一条线确认 loss 能下降后再改成 5-shot 看精度提升。直接上 20-way 容易在调参时把自己搞晕因为每次损失下降的幅度都更小不好判断是模型问题还是任务太难。3. 核心模块拆解一个 Episode 从采样到返回损失3.1 原型损失函数三个步骤完成一次学习整个项目的灵魂是prototypical_loss这个函数它把“计算原型 → 计算距离 → 计算损失”三步串在一起。我拆给你看。def prototypical_loss(proto_net, support_x, support_y, query_x, query_y, n_way, n_shot, n_query): 原型网络的单次 episode 损失 Args: proto_net: 嵌入网络输入图像输出 embedding 向量 support_x: 支持集图像shape [n_way * n_shot, C, H, W] support_y: 支持集标签shape [n_way * n_shot] query_x: 查询集图像shape [n_way * n_query, C, H, W] query_y: 查询集标签shape [n_way * n_query] Returns: loss: 交叉熵损失 acc: 查询集准确率 # 支持集和查询集各自过嵌入网络 support_emb proto_net(support_x) # [n_way * n_shot, embedding_dim] query_emb proto_net(query_x) # [n_way * n_query, embedding_dim] # 第一步按类别分组计算每个类的原型 prototypes [] for i in range(n_way): # 找出支持集里属于第 i 类的样本下标 indices (support_y i).nonzero(as_tupleFalse).squeeze(1) class_emb support_emb[indices] # [n_shot, embedding_dim] proto class_emb.mean(dim0) # [embedding_dim] prototypes.append(proto) prototypes torch.stack(prototypes) # [n_way, embedding_dim] # 第二步查询样本与所有原型计算欧氏距离 dist torch.cdist(query_emb, prototypes) # [n_way * n_query, n_way] # 第三步距离取负过 softmax得到分类概率 logits -dist loss F.cross_entropy(logits, query_y) acc (logits.argmax(dim1) query_y).float().mean().item() return loss, acc代码的巧妙之处在于torch.cdist一行算完了所有查询点和所有原型之间的距离矩阵不需要手写双层循环。logits取负是因为“距离越小概率越大”而交叉熵自带 softmax所以不用单独再套一层。有个细节需要注意support_y的标签必须是从 0 开始的连续整数。如果每个 episode 的类别编号是从全局类别 id 里直接拿来的比如 17、89、203那range(n_way)就对不上索引原型会被算到错误的分组里。常见做法是在采样器里做一次“重编号”把全局类别 id 映射成 0 到 n_way-1 的局部 id。3.2 嵌入网络用 Conv4 把图像压成 64 维向量原型网络对嵌入网络的要求是“结构简单、不容易过拟合”。因为 Few-shot 场景本身数据量很少用 ResNet50 这种大网络很容易在支持集上直接记住样本反而泛化更差。这个项目默认用的是 Conv4四层卷积加一个全连接层。class Conv4(nn.Module): 轻量嵌入网络4 层卷积自适应池化全连接输出输出维度是 embedding_dim通常是 64 def __init__(self, in_channels1, hidden_dim64, embedding_dim64): super().__init__() self.features nn.Sequential( nn.Conv2d(in_channels, hidden_dim, 3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, 3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, 3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, 3, padding1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d(1), # 把空间尺寸压成 1x1任何输入尺寸都通用 ) self.fc nn.Linear(hidden_dim, embedding_dim) def forward(self, x): x self.features(x) # [B, hidden_dim, 1, 1] x x.view(x.size(0), -1) # [B, hidden_dim] return self.fc(x) # [B, embedding_dim]注意AdaptiveAvgPool2d(1)这个设计。很多初版代码会直接用x.view(x.size(0), -1)然后接全连接层这要求输入图片尺寸在训练和测试时必须完全一致。换成自适应池化后不管输入是 84x84 还是 128x128进全连接层之前都被压成 1x1换数据集的时候少踩一个尺寸不匹配的坑。BatchNorm 在这里不是可有可无的。原型计算依赖嵌入空间的稳定性BatchNorm 能抑制不同 episode 之间因为随机采样带来的特征分布抖动。我试过去掉 BatchNorm训练 loss 还是能降但验证准确率波动明显变大尤其在 1-shot 场景下差 3 到 5 个点都很正常。3.3 训练循环与采样器完整的最小实现前面两个模块讲的是“单次推理”训练还缺一个关键的拼图——Episode 采样器。它的职责是每次生成一个全新的小任务。def sample_episode(dataset, n_way, n_shot, n_query): 从数据集中采样一个 episode 随机选 n_way 个类别每类抽 n_shot 张支持集、n_query 张查询集 # dataset.classes 是全局类别列表 # dataset.class_to_indices 是 dict: 类别名 - 样本索引列表 selected_classes random.sample(dataset.classes, n_way) support_x, support_y [], [] query_x, query_y [], [] for local_id, cls_name in enumerate(selected_classes): indices dataset.class_to_indices[cls_name] random.shuffle(indices) # 前 n_shot 张作支持集后 n_query 张作查询集 support_idx indices[:n_shot] query_idx indices[n_shot:n_shot n_query] support_x.extend(dataset.get_images(support_idx)) support_y.extend([local_id] * n_shot) query_x.extend(dataset.get_images(query_idx)) query_y.extend([local_id] * n_query) return (torch.stack(support_x), torch.LongTensor(support_y), torch.stack(query_x), torch.LongTensor(query_y))这段代码里最容易出错的是get_images和class_to_indices的配合。class_to_indices存的是图片索引而不是图片本身这样采样时只操作整数列表内存开销小真正读图发生在get_images里每张图只读一次避免把所有图片都塞进内存。训练循环就很常规了——采样一个 episode、算损失、反向传播、更新参数重复到 epoch 结束。def train_epoch(proto_net, optimizer, loader, n_way, n_shot, n_query, device): proto_net.train() total_loss, total_acc 0.0, 0.0 num_batches 0 for batch in loader: support_x, support_y, query_x, query_y batch support_x support_x.to(device) support_y support_y.to(device) query_x query_x.to(device) query_y query_y.to(device) loss, acc prototypical_loss( proto_net, support_x, support_y, query_x, query_y, n_way, n_shot, n_query ) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() total_acc acc num_batches 1 avg_loss total_loss / num_batches avg_acc total_acc / num_batches return avg_loss, avg_acc整个实现最让我感慨的是原型网络的核心逻辑简单到没有秘密难的全在采样器和预处理这两个外围环节。很多人跑不出论文里的指标问题往往不在模型代码而在 episode 采样是否均匀、预处理是否与训练一致。4. 运行与常见问题排查从环境搭建到训练曲线4.1 环境准备与启动命令这个项目依赖很少PyTorch 2.x 加上 torchvision、tqdm、pyyaml 就够。我建议用 conda 单独建一个环境避免把其他项目的依赖搞乱。conda create -n proto python3.10 conda activate proto pip install torch torchvision tqdm pyyaml python train.py --config configs/default.yamldefault.yaml里面约定了 n_way、n_shot、n_query、学习率和训练轮数。第一次跑的时候不要改任何参数先看看 loss 和准确率的变化趋势。正常情况是 loss 从接近 ln(n_way)比如 5-way 就是 1.6 左右开始下降验证准确率在几十个 epoch 内稳步上升。训练过程中会在logs/目录下生成 CSV 格式的训练记录用任何表格软件都能打开。我习惯每隔 5 个 epoch 打印一次验证准确率看得比 loss 更直观。4.2 踩坑记录一把欧氏距离换成余弦就废了现象模型 loss 能降但验证准确率始终卡在 1/n_way 附近和随机猜测差不多。原因很多初学者上手第一件事就是把torch.cdist换成余弦相似度觉得“语义相似度用余弦更合理”。但原型网络的核心前提是“原型类内均值”这个均值在欧氏空间里才是最优代表。换成余弦后嵌入向量的模长会干扰均值方向类别均值不再是方向上的中心距离比较也就失去了意义。解决保持欧氏距离如果确实想用余弦必须先对嵌入向量做 L2 归一化再算欧氏距离。我常用的是F.normalize(embed, p2, dim-1)之后再喂给torch.cdist这时欧氏距离和余弦距离在数学上等价均值的意义也能保留。4.3 踩坑记录二训练集做的预处理验证集必须原样复刻现象训练 loss 平滑下降但每跑一个 epoch 验证准确率都剧烈抖动甚至比训练低十几个点。原因训练时用了随机裁剪、随机翻转、颜色抖动等数据增强验证 / 测试时却只做了 resize两边的数据分布不一致。嵌入网络在增强后的分布上学到的特征在验证集的干净图上不一定起作用。解决把训练和验证的预处理写成一个共享函数唯一的区别是训练时多一步随机增强。更保险的做法是固定验证集的随机种子保证每次验证看到的图片顺序一致这样准确率曲线不会因为采样抖动神一阵鬼一阵。4.4 踩坑记录三Windows 下 DataLoader 卡死与随机种子玄学现象程序启动后卡在第一个 batchCPU 占用率忽高忽低训练进度条一动不动。原因Windows 上 PyTorch 的 DataLoader 多进程使用 spawn 模式如果训练代码没有被if __name__ __main__:保护子进程会重复执行模块级代码导致死锁。另外num_workers设置过大且每个 worker 都各自初始化随机数生成器会让不同运行的结果差异巨大。解决在train.py末尾加保护把num_workers调成 0 先跑通再逐步加。至于随机种子我的做法是训练入口固定一次全局 seed然后把“采样器是否 shuffle”作为唯一的随机性来源这样每次运行至少能保证结果可复现出了问题能回头查。5. 二次开发把这个实现改成你自己的小样本任务5.1 迁移到自定义数据集只需改采样器和预处理很多朋友拿到这份代码第一个念头是“我能直接用吗”。答案是能但要把数据接口换一下。原项目的数据集类做了三件事扫描目录、维护class_to_indices映射、实现get_images。你换成自己的数据集时只需要保证这三件事的接口签名不变模型和训练循环一行都不用改。class CustomDataset: def __init__(self, root_dir, transformNone): self.classes [] self.class_to_indices {} self.image_paths [] # 扫描 root_dir 下每个子文件夹作为一个类别 for class_id, class_name in enumerate(os.listdir(root_dir)): class_dir os.path.join(root_dir, class_name) indices [] for img_name in os.listdir(class_dir): self.image_paths.append(os.path.join(class_dir, img_name)) indices.append(len(self.image_paths) - 1) self.class_to_indices[class_name] indices self.classes.append(class_name) self.transform transform def get_images(self, indices): images [] for idx in indices: img Image.open(self.image_paths[idx]).convert(RGB) if self.transform: img self.transform(img) images.append(img) return images这里有个容易被忽略的点类别扫描顺序要固定。Python 的os.listdir在不同系统上返回顺序不稳定如果classes列表顺序每次启动都变selected_classes的采样结果也会变测试时切到新数据集对不上旧结果。用sorted(os.listdir(...))固定顺序算是我的一个血泪经验。5.2 替换 Backbone 与距离度量从 Conv4 到 ResNetConv4 在小数据集上很稳但如果你手里的数据量稍微大一点、图像分辨率高一点想提点最直接的改法是换 Backbone。import torchvision.models as models def build_encoder(backboneresnet18, embedding_dim64): if backbone resnet18: base models.resnet18(weightsNone) # 去掉最后的全连接层只保留特征提取部分 base.fc nn.Identity() # 接一个线性层把特征压缩到 64 维 encoder nn.Sequential( base, nn.Linear(512, embedding_dim) ) return encoder换了 Backbone 之后要同步调学习率。ResNet 的初始特征提取能力比 Conv4 强但微调时学习率太大会破坏预训练权重。我一般把 backbone 的学习率设为新的线性层的 0.1 倍用参数分组实现。距离度量方面把torch.cdist换成余弦距离也只需要两行# 欧氏距离版本dist torch.cdist(query_emb, prototypes) # 余弦版本先 L2 归一化再算欧氏距离 query_emb F.normalize(query_emb, p2, dim-1) prototypes F.normalize(prototypes, p2, dim-1) dist torch.cdist(query_emb, prototypes)不过要记住改了距离度量就等于改了模型的核心假设一定要重新验证均值作为原型是否还合理。从那以后我每次换数据集或换 Backbone都会先用一个小模型、小学习率把过拟合跑通确认 loss 在动、准确率能涨然后再放大规模去调终极指标。这套流程帮我省掉了不少“看着 loss 下降但验证就是上不去”的玄学调试时间希望帮到你。本文还有配套的精品资源点击获取

PyTorch原型网络实战：少样本分类的轻量级方案与代码详解

相关推荐

2S电池组主动均衡系统：ATmega644+MP2642GR-Z设计实战

组合总和（LeetCode 39）题解：排序 + 剪枝 + 回溯两种写法的完整剖析

instagram-private-api 评论响应用户结构体 `MediaRepositoryCommentResponseUser` 全字段解析与实战指南

更多资讯

四小时学会OpenCV+Qt：图像查看器开发全流程

Jenkins+SVN+Unity自动打包搭建指南：从环境配置到常见坑

Office 365与AD域同步实战：从部署到排错全攻略

HugAgentOS接入飞书/钉钉/企微：如何打造看得到群聊上下文的AI机器人

B站批量取关完整指南：BiliBiliToolPro 定时清理关注列表

draw.io 桌面版完整安装指南：免费绘图工具覆盖流程图、UML 与架构图

今日更新

Cephadm实战指南：Linux系统工程师的生产级Ceph集群部署与排障

不想再被 Token 额度卡脖子？把免费 API 接入 WorkBuddy 的完整教程

ITK中SpatialObject几何体系解析：坐标系、变换与管道集成

本周精选

交通拥堵预测毕设：基于Python+LSTM的车辆流量预测系统全流程解析

MATLAB/Simulink搭建10机39节点电力系统暂态稳定仿真实战

无人机航拍三维重建全流程：从SfM到网格生成的避坑指南

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？