简介:这是一个基于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 或 5 | 1-shot 是难度上限 |
| n_query | 每类查询集样本数 | 15 或 16 | 影响梯度估计稳定性 |
| meta_batch_size | 每次迭代的 episode 数 | 1 | 显存有限时固定为 1 |
| lr | 学习率 | 1e-3 到 1e-4 | 用小网络时 1e-3 起步 |
| embedding_dim | 嵌入向量维度 | 64 | 维度太低表达能力受限 |
我一般会先用 n_way=5、n_shot=1、n_query=15 跑通一条线,确认 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_tuple=False).squeeze(1) class_emb = support_emb[indices] # [n_shot, embedding_dim] proto = class_emb.mean(dim=0) # [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(dim=1) == 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_channels=1, hidden_dim=64, embedding_dim=64): super().__init__() self.features = nn.Sequential( nn.Conv2d(in_channels, hidden_dim, 3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), nn.BatchNorm2d(hidden_dim), nn.ReLU(inplace=True), 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 python=3.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, p=2, dim=-1)之后再喂给torch.cdist,这时欧氏距离和余弦距离在数学上等价,均值的意义也能保留。
4.3 踩坑记录二:训练集做的预处理,验证集必须原样复刻
现象:训练 loss 平滑下降,但每跑一个 epoch 验证准确率都剧烈抖动,甚至比训练低十几个点。
原因:训练时用了随机裁剪、随机翻转、颜色抖动等数据增强,验证 / 测试时却只做了 resize,两边的数据分布不一致。嵌入网络在增强后的分布上学到的特征,在验证集的干净图上不一定起作用。
解决:把训练和验证的预处理写成一个共享函数,唯一的区别是训练时多一步随机增强。更保险的做法是固定验证集的随机种子,保证每次验证看到的图片顺序一致,这样准确率曲线不会因为采样抖动神一阵鬼一阵。
4.4 踩坑记录三:Windows 下 DataLoader 卡死与随机种子玄学
现象:程序启动后卡在第一个 batch,CPU 占用率忽高忽低,训练进度条一动不动。
原因: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, transform=None): 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 到 ResNet
Conv4 在小数据集上很稳,但如果你手里的数据量稍微大一点、图像分辨率高一点,想提点,最直接的改法是换 Backbone。
import torchvision.models as models def build_encoder(backbone='resnet18', embedding_dim=64): if backbone == 'resnet18': base = models.resnet18(weights=None) # 去掉最后的全连接层,只保留特征提取部分 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, p=2, dim=-1) prototypes = F.normalize(prototypes, p=2, dim=-1) dist = torch.cdist(query_emb, prototypes)不过要记住,改了距离度量就等于改了模型的核心假设,一定要重新验证均值作为原型是否还合理。从那以后我每次换数据集或换 Backbone,都会先用一个小模型、小学习率把过拟合跑通,确认 loss 在动、准确率能涨,然后再放大规模去调终极指标。这套流程帮我省掉了不少“看着 loss 下降但验证就是上不去”的玄学调试时间,希望帮到你。
本文还有配套的精品资源,点击获取