PyTorch实现原型网络：小样本图像分类实战指南

简介：原型网络（Prototypical Networks）是少样本学习领域的经典方法，该PyTorch实现对应论文arXiv:1703.05175，面向机器学习研究者、算法工程师及相关方向学生，帮助他们快速复现基准模型并开展小样本分类实验。资源基于Omniglot数据集提供完整工程代码，涵盖数据读取、类别原型计算、距离度量、损失函数与训练测试流程，适合具备一定深度学习基础、希望复现或对比原型网络性能的开发者。压缩包共12个文件，容量仅135KB；核心为7个Python脚本，分别处理参数解析、批次采样、模型搭建、损失计算和训练入口，代码结构清晰便于二次开发；另含2张示意图、README说明、License与gitignore，可快速了解运行方式与许可条款。目前已有920人学习。通过该资源，读者可掌握原型网络在Omniglot上的训练细节，理解episode机制与原型更新逻辑，并可将代码迁移至自定义少样本数据集，是入门小样本学习领域的高质量参考资料。

1. 为什么「只有几个标签」的分类任务，比数据增强更适合原型网络

去年做一个模拟项目X的缺陷分类，正样本只有十几个，负样本倒是拉了几百张。常规思路是拿预训练模型微调，结果不是过拟合就是把负样本全认成正样本，翻车翻得我怀疑人生。后来把方案换成这个 PyTorch 实现的原型网络项目，反而几分钟就把基线跑到了能用的水平。这个项目解决的，正是机器学习里被问得最多的问题：每个类只给 5 张甚至 1 张标注图，分类器还能不能学会。

它的思路其实很朴素：每个类别在特征空间里有个「原型」，也就是同类样本特征的平均；新样本离哪个原型近就归哪类。这和你记忆里「聚类」有点像，但它把原型和距离都放进了一个端到端训练的神经网络里。适合谁？像我这样做小样本图像分类、低成本冷启动的 Python 开发者。你要是刚入门 few-shot learning，这个项目也是极好的第一份可跑源码。

2. 项目结构与运行环境：小样本仓库怎么组织才能不跑偏

2.1 目录布局：模型、数据、脚本各司其职

这份源码的整体结构非常朴素，没有把代码揉成一大坨，直接按模型、数据、训练入口、评估入口分开放。这也是我见过比较舒服的小样本项目组织方式：

prototypical/ ├── models/ │ ├── __init__.py │ └── prototypical_net.py # 原型网络核心实现 ├── data/ │ └── few_shot_dataset.py # 按 episode 采样的数据加载 ├── configs/ │ └── default.yaml # 超参数统一入口 ├── utils/ │ ├── metrics.py # 准确率、置信区间 │ └── seed.py # 随机种子管理 ├── train.py # 训练入口 └── evaluate.py # 评估入口

重点看两个文件：prototypical_net.py是模型本体，few_shot_dataset.py负责生成 episode。小样本学习和普通分类最大的区别就在这里——普通数据集是「一批图片 + 一批标签」，小样本数据集是「一个任务 + 支持集 + 查询集」。刚开始读这个项目，建议先打开few_shot_dataset.py看明白 episode 怎么生成，再看模型，否则看训练代码时会一头雾水。

2.2 环境准备：从零到第一个 episode

我先说结论：这个项目对环境的需求非常简单，不挑 GPU，CPU 也能跑通小数据集上的 baseline，只是慢一点。我一般会新建一个干净的虚拟环境：

conda create -n proto python=3.10 -y conda activate proto pip install torch torchvision torchaudio pyyaml

如果你的机器有 GPU，记得装对应 CUDA 版本的 PyTorch，别用默认的 CPU 版本，否则训练速度会劝退你。这里没有引入复杂的依赖，项目本身的代码几乎只依赖 PyTorch、NumPy 和 PyYAML。pyyaml用来读configs/default.yaml，如果你不想用配置文件，直接改成 argparse 也是可以的。

启动训练之前，建议先跑一个最小配置：n_way=5，n_shot=5，n_query=15，这个配置在大多数小样本数据集上都能得到一个比较稳的基准。我第一次跑的时候用了n_way=2，结果损失曲线像个心电图，后面在避坑章节我会专门讲这个。

2.3 数据准备：支持集和查询集怎么划分

这个项目不要求你把数据整理成特殊的二进制格式，普通的图片文件夹就能用。数据目录大概是：

data/ ├── train/ │ ├── class_001/ │ │ ├── sample_01.jpg │ │ └── ... │ └── class_002/ └── val/ ├── class_001/ └── class_002/

每个子文件夹是一个类别，文件夹名就是类别标签。项目启动时会把目录里的类别列出来，然后按 episode 采样。这里容易踩的一个坑是：训练集和验证集的类别必须是互斥的。少样本学习评估的是「模型有没有学会按原型分类」，而不是「它有没有记住训练类」，所以验证集类别不能出现在训练集里。如果有人用传统的train_test_split随机切图片，那评估出来的准确率一定是虚高的。

3. 核心模块拆解：一个 episode 从采样到损失回传的完整链路

3.1 特征提取器：为什么这里不直接接全连接分类层

原型网络的骨干部分和普通分类网络没什么区别，但最后面不接全局池化 + 全连接分类层，因为我们要的不是类别 logits，而是能把同类样本聚在一起的特征向量。项目里默认的骨干是一个四层卷积网络，每层是「卷积 + 批归一化 + ReLU + 最大池化」，这个结构在少样本场景里比直接拿大模型更稳：

import torch.nn as nn def conv_block(in_channels, out_channels): return nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.MaxPool2d(2), ) class Conv64F(nn.Module): def __init__(self, input_channels=3): super().__init__() self.encoder = nn.Sequential( conv_block(input_channels, 64), conv_block(64, 64), conv_block(64, 64), conv_block(64, 64), ) def forward(self, x): # 输入: [batch_size, C, H, W] out = self.encoder(x) # 输出: [batch_size, 64 * H_ * W_] return out.view(out.size(0), -1)

这里的参数含义要说明一下：input_channels是输入图像的通道数，灰度图设为 1，RGB 图设为 3。每次卷积后特征图尺寸减半，四层之后一张84x84的输入图会变成5x5大小，特征维数是64*5*5=1600。这个 1600 维的向量就是后续计算原型的输入。

为什么小样本训练反而不太推荐一上来就上 ResNet？因为每类只有几张图，骨干网络参数太多，很容易把所有支持集图像记忆住，然后查询集一进来就崩。四层卷积网络虽然结构简单，但特征维度适中，原型计算时方差更可控，这是项目原作者反复验证过的选择。

3.2 原型计算与距离度量：欧氏距离才是灵魂

拿到特征之后，计算原型只是一个求均值操作。假设一个 episode 的配置是 5 个类别、每类 5 张支持图，那么支持集特征形状是[25, 1600]，按类别分组后求均值，得到 5 个原型向量：

import torch def get_prototypes(support_features, n_way, n_shot): # support_features: [n_way * n_shot, feature_dim] support_features = support_features.view(n_way, n_shot, -1) # 返回: [n_way, feature_dim] return support_features.mean(dim=1)

这里有一个很多人会忽略的细节：mean(dim=1)是沿着 shot 维度取平均，也就是把每个类别的 n_shot 个样本特征逐元素相加后除以 n_shot。这也意味着原型向量和每张图处于同一个特征空间，距离计算才有意义。

查询集的分类概率，用负的平方欧氏距离经过 softmax 得到：

def prototypical_loss(query_features, prototypes, query_labels, n_query): # query_features: [n_query_total, feature_dim] # prototypes: [n_way, feature_dim] dist = torch.cdist(query_features, prototypes, p=2).pow(2) logits = -dist log_p = logits.log_softmax(dim=1) # 收集每个查询点的真实类别下标 target = query_labels.view(-1) loss = torch.nn.functional.nll_loss(log_p, target) return loss

torch.cdist这个函数是欧氏距离的关键，它直接计算查询特征和所有原型之间的两两距离。pow(2)表示取平方，因为后面要送进 log_softmax，平方能让距离差异更明显。这里不建议换成余弦相似度，因为余弦只看方向不看长度，而原型均值这一步会把向量的模长信息混进去，实验里经常出现「准确率看似不低、但置信区间特别大」的诡异现象。

3.3 episode 采样：少样本学习的「任务」单位

普通数据集的 DataLoader 返回一个 batch，小样本项目的 DataLoader 返回一个 episode。这个 episode 里包含「支持集」和「查询集」两部分，项目里的采样逻辑可以单独抽出来看：

import numpy as np class FewShotEpisodeSampler: def __init__(self, labels, n_way, n_shot, n_query, episodes): self.labels = np.array(labels) self.classes = np.unique(labels) self.n_way = n_way self.n_shot = n_shot self.n_query = n_query self.episodes = episodes def __iter__(self): for _ in range(self.episodes): # 1. 从所有类别中选出 n_way 个类 episode_classes = np.random.choice(self.classes, self.n_way, replace=False) support_indices = [] query_indices = [] for cls in episode_classes: idx = np.where(self.labels == cls)[0] # 2. 每个类先抽 n_shot 个作为支持 selected = np.random.choice(idx, self.n_shot + self.n_query, replace=False) support_indices.extend(selected[:self.n_shot]) query_indices.extend(selected[self.n_shot:]) yield support_indices, query_indices, episode_classes

这个采样器里最重要的参数是replace=False。许多少样本项目准确率虚高，就是因为采样时没有做无放回抽取，同一张图既出现在支持集又出现在查询集，模型等于「开卷考试」，验证分数一点参考价值都没有。这里先抽n_shot + n_query，再前n_shot个做支持、剩下的做查询，是为了保证同一个类别两个集合的图像不重叠。

3.4 训练损失与梯度回传：小 batch 更新背后的数学直觉

整个训练流程可以压缩成一个 forward 函数：支持集过骨干网络得到原型，查询集过骨干网络得到查询特征，然后计算负对数似然损失。这个损失的意义很直接——我们希望查询点被分配到正确原型的概率尽可能接近 1。

logits = -dist loss = logits.log_softmax(dim=1).nll_loss(target)

很多第一次读源码的人会纳闷：为什么不直接用CrossEntropyLoss？其实CrossEntropyLoss内部就是log_softmax + nll_loss，这里分开写是为了让代码更透明。梯度回传方向和普通分类网络完全一致，loss.backward()会把梯度一路传回骨干网络。

小样本训练和普通训练最大的不同是每个 batch 里的样本数很少，比如n_way=5, n_shot=5, n_query=15，一个 episode 只有 40 张图，但类别却有 5 个，所以每个类别的梯度会非常「尖锐」。我一般建议学习率从1e-3起步，如果发现 loss 在第一个 episode 就快速跌到零点几，然后开始震荡，大概率是学习率偏大，后面避坑章节会展开。

4. 训练与评估：从命令行到准确率的完整闭环

4.1 训练主循环：一个 epoch 到底在迭代什么

这个项目的训练脚本逻辑清晰，和普通 PyTorch 训练循环相比，只是把「batch」换成了「episode」：

from utils.seed import set_seed def train_one_epoch(model, sampler, optimizer, backbone, n_way, n_shot, n_query): model.train() total_loss = 0.0 for support_idx, query_idx, cls in sampler: support_images = load_images(support_idx) query_images = load_images(query_idx) optimizer.zero_grad() support_feat = backbone(support_images) query_feat = backbone(query_images) prototypes = get_prototypes(support_feat, n_way, n_shot) loss = prototypical_loss(query_feat, prototypes, query_labels, n_query) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(sampler)

注意这里的load_images只是一个示意函数，实际项目里会从内存或者磁盘批量加载并做归一化。每次迭代只有一个 episode，也就是 5 个类、每类 4 张左右图。因此一个 epoch 不再是「所有数据过了一遍」，而是「采样器生成的所有 episode 过了一遍」。默认配置下 1000 个 episode 为一轮，我习惯跑 60 轮，因为小样本数据总量不大，轮数可以适当多一些。

4.2 评估流程：固定种子下的多次采样取均值

评估脚本evaluate.py的核心逻辑是固定随机种子，在验证集上反复生成新的 episode，算平均准确率。为什么要固定种子？因为小样本评估结果方差太大，同一个模型换一批随机 episode，准确率可能差 5 个百分点。不固定种子，你根本没法比较两个模型的优劣。

def evaluate(model, val_labels, n_way=5, n_shot=5, n_query=15, episodes=1000, seed=42): set_seed(seed) sampler = FewShotEpisodeSampler(val_labels, n_way, n_shot, n_query, episodes) correct = 0 total = 0 with torch.no_grad(): for support_idx, query_idx, cls in sampler: support_feat = model(load_images(support_idx)) query_feat = model(load_images(query_idx)) prototypes = get_prototypes(support_feat, n_way, n_shot) dist = torch.cdist(query_feat, prototypes, p=2) preds = dist.argmin(dim=1) correct += (preds == query_labels).sum().item() total += query_labels.size(0) return correct / total

这段代码里唯一和分类网络不同的地方是，最后没有 softmax，而是直接取距离最小的原型作为预测类别。因为softmax不改变argmin的结果，所以评估时去掉它还能省一点计算。我通常在验证时同时计算准确率和 95% 置信区间，只报一个均值很容易被随机性欺骗。

4.3 指标与日志：准确率之外还要看什么

小样本学习的评估指标以准确率为主，但只看准确率会掉坑。下表是我在这个项目里常用的一组指标：

训练时我一般同时记录 loss 和当前 episode 的准确率，但真正决定是否早停的是验证集 mean accuracy。如果训练 loss 一直降，验证准确率却掉到随机水平，那大概率是骨干网络遇到数值问题或者数据泄露，不是模型能力问题。

5. 避坑指南：小样本训练最常见的五个「玄学」现场

5.1 支持集和查询集信息泄露

现象：训练准确率飙到 98%，但验证准确率只有 70%。原因：采样时没有区分同一个类的支持集和查询集，或者验证集和训练集类别有重叠。解决：严格检查FewShotEpisodeSampler的replace=False逻辑，同时在数据划分时保证训练类别和验证类别完全互斥。我血的教训是，第一次实现时从每个类别的所有图片里随机抽支持集，查询集也从同类别里随机抽，两者可能抽到同一张图片。后来加了「先抽n_shot + n_query，再按位置切分」的逻辑才修掉。

5.2 距离度量选错

现象：训练时 loss 很低，但验证准确率忽高忽低，换一个随机种子结果差很多。原因：用了余弦相似度或者曼哈顿距离。原型网络的原生假设是欧氏距离，因为原型是特征均值，均值操作天然和欧氏空间绑定；余弦相似度只关心方向，忽略向量的长度信息，导致原型的位置表达失效。解决：用torch.cdist(features, prototypes, p=2)计算平方欧氏距离，不要自行实现余弦。

5.3 n_way 设置太小

现象：损失曲线像心电图，准确率长期停在 0.5 附近。原因：n_way=2时，模型只需要区分两个类，梯度信号过于单一，而且每次 episode 只有两个原型，特征空间里大量区域没有约束，训练很难稳定。解决：把n_way设为 5 或者更高的类别数，保证每次 episode 有足够的类别间对比信号。这不是什么玄学，本质是让同一批特征要同时分开多个类，原型之间的相对位置才学得清楚。

5.4 随机种子不固定

现象：同一份代码、同一个超参数，跑两次结果完全不一样，甚至后一次准确率比前一次低 8 个百分点。原因：小样本评估依赖随机 episode 采样，数据加载顺序、数据增强、模型初始化都会引入随机性。解决：在训练和评估前后都调用set_seed(seed)，包括 numpy、torch、random 三个库；如果要严格复现，还要关掉 cudnn 的自动 benchmark。我一般把种子写进配置文件，模型文件命名里带上种子号。

5.5 学习率过大导致原型「漂移」

现象：第一个 episode 的 loss 只有 0.2，之后十几个 episode 一路冲到 3.0，再也降不回来。原因：小样本每个 episode 样本量太少了，学习率大时梯度一步就把骨干网络参数推到特征分布扭曲的位置，原型均值跟着漂移。解决：把学习率降到1e-3以下，必要时先跑 5 个 episode 观察 loss；或者加一个学习率预热，前几个 episode 只用非常小的学习率热身。

6. 二次开发：把原型网络改造成自己的小样本分类器

6.1 替换骨干网络：从四层卷积到预训练模型

项目默认的Conv64F适合快速验证思路，但如果你的数据集不是那种「每类几十张」的玩具规模，换成更大的骨干网络能直接提升准确率。改造点只有一个：把Conv64F换成输出特征向量固定维度的网络，然后保留原型计算部分不动。

import torchvision.models as models class ResNetBackbone(nn.Module): def __init__(self, output_dim=512, pretrained=True): super().__init__() resnet = models.resnet18(pretrained=pretrained) self.features = nn.Sequential(*list(resnet.children())[:-1]) self.fc = nn.Linear(512, output_dim) def forward(self, x): feat = self.features(x) feat = feat.view(feat.size(0), -1) return self.fc(feat)

换完之后注意两个地方：输入图像尺寸要改成你预训练模型要求的大小，比如 ResNet 系列一般是224x224；骨干网络如果冻结了前几层，要确保requires_grad设置正确。经验是，初始阶段冻结layer1之前的参数，只训练后面的层和输出层，效果会比全量微调更稳。

6.2 加入针对少样本的数据增强

原始原型网络训练时通常不强求复杂增强，但真实项目里每个类的样本实在太少，不做增强很容易让原型向量偏向某一张图的噪声。我后来给这个项目加了一个轻量增强流程：对支持集做RandomCrop和ColorJitter，对查询集只做中心裁剪和归一化。为什么支持集和查询集增强策略不一样？因为查询集模拟的是真实推理时的输入，过度增强会让训练和推理输入分布不一致。

train_transform = transforms.Compose([ transforms.RandomResizedCrop(84), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.3, 0.3, 0.3), transforms.ToTensor(), transforms.Normalize(mean, std), ])

注意增强强度的控制。少样本场景下增强太猛，比如随机擦除、极亮颜色抖动，会把本来就不多的类别特征彻底扭曲，原型均值直接失去代表性。我的经验是先把增强强度设为普通分类任务的一半，再根据验证准确率微调。

6.3 自定义 episode：按难度采样的进阶玩法

如果你想在这个项目上做研究或者解决实际问题，一个很有效的小改动是修改FewShotEpisodeSampler，让它不只是均匀随机选类别，而是按类别间相似度挑难度更高的任务。比如把特征向量直接用于计算类别相似度，优先抽相似度高的类别组成 episode，这样模型被迫学会更精细的原型边界。

def __iter__(self): class_feats = compute_class_prototypes(self.train_data) similarity = cosine_similarity(class_feats) for _ in range(self.episodes): anchor = np.random.choice(self.classes) similar_classes = np.argsort(similarity[anchor])[-self.n_way:] episode_classes = np.random.choice(similar_classes, self.n_way, replace=False) # 后续采样逻辑不变

这个改动不会破坏原有模型结构，只是让训练任务变得更难。从结果来看，验证集准确率可能不会立刻变高，但模型的鲁棒性会明显提升。从那以后我每次做小样本相关的改造，都会强制跑完整链路：先固定种子训练一轮，看 loss 是否收敛，再换 1000 个 episode 评估置信区间，最后才敢把模型交给下游业务。这个流程听起来繁琐，但每一次都能提前暴露问题，希望帮到你。

本文还有配套的精品资源，点击获取