ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

小样本学习下的北极熊分割识别:原型网络与PyTorch实战

小样本学习下的北极熊分割识别:原型网络与PyTorch实战 简介2018年第八届华为杯竞赛参赛项目资料包主题是“基于小样本学习的自然场景北极熊高效分割识别系统”面向计算机视觉学习者、竞赛选手以及小样本算法研究人员。项目核心是利用有限北极熊图像完成高效分割与识别方案中涵盖小样本学习常用方法元学习、原型学习、迁移学习、语义/实例分割模型Mask R-CNN、YOLACT、轻量化网络MobileNet、ShuffleNet以及数据增强、正则化、集成学习等优化技术适合用于算法复现、赛题复盘或思路扩展。压缩包共295个文件、18.56MB以hpp112个和h55个头文件为主另有cpp源码、lib/dll链接库、obj中间文件、jpg/png/bmp图像样本及sln/vcxproj工程文件目录结构清晰既能查看算法实现也可尝试重新编译运行。目前已有96人学习下载。包内还包含pdf说明、md笔记、日志与工程缓存文件能够帮助读者理解数据划分、模型评估指标IoU、精确率、召回率、F1在真实竞赛中的落地方式。1. 当训练集只有几十张图时北极熊分割识别怎么做小样本学习Few-Shot Learning与分割识别系统的组合在 2018 年的竞赛场景里是一个很前卫的技术方向参赛者面对的北极熊影像来源可能是无人机航拍或野外红外相机标注样本往往只有几十张而自然场景里熊的毛色与雪地、岩石、水面高度相近背景杂乱、光照多变。传统语义分割模型在数据充足时能轻松达到 80% 以上的 mIoU但样本量骤减后模型不是过拟合就是学不到熊的轮廓特征。这个标题背后的真实需求就是要在少量标注样本下让分割识别系统仍然具备可用的精度和鲁棒性。适合的读者是正在做算法竞赛、工业视觉小样本落地的工程师以及想在小样本方向入门的研究生。2. 先定方案再写代码小样本分割的架构选型与数据划分2.1 把“分割识别”拆成像素级的 N-way K-shot 问题小样本学习的经典设定是 N-way K-shot支持集Support Set中每个类别提供 K 张带标注的样本模型在这 K 张样本的“指导”下对查询集Query Set中未标注的图像做预测。放在北极熊分割识别场景下任务被改写为像素级的二分类前景北极熊与背景。这里的 N 通常取 1K 视标注成本取 1 到 8 之间的整数。竞赛选手的实际做法是用少量标注图像生成支持集其余图像作为查询集模型必须从支持集中“提炼”出北极熊的外观特征再把这个特征迁移到查询集的每一个像素上。我在处理这类任务时的经验是不要一上来就堆模型先把数据划分做对。常见做法是把训练数据拆成三份基类数据集、支持集、查询集。基类数据集可以来自公开数据集中与北极熊外观相似的动物类别比如狗、熊科其他物种用于预训练特征提取器。支持集与查询集则来自北极熊的真实影像且二者必须保证视频帧或拍摄位置的独立性否则会带来数据泄漏这一点会在后面的避坑章节详谈。2.2 为什么选 ResNet-50 原型对齐而不是直接微调 Mask R-CNN2018 年前后Mask R-CNN 已经是实例分割的默认选项很多参赛者第一反应就是拿它来微调。但小样本设定下直接微调 Mask R-CNN 存在两个致命问题一是 RPNRegion Proposal Network层需要大量正负样本才能训练稳定几十张图远不够二是检测头只有少量样本驱动时分类分支极易过拟合熊的边界框稍微偏移一点分割掩膜就跟着崩掉。更稳妥的选择是基于原型网络Prototype Network思路的编码器-原型器结构。核心思想是用基类数据预训练一个特征提取器让它在语义层面能区分“什么是有生命的物体”与“什么是背景”然后在推理阶段把支持集标注的掩膜区域内的特征做平均池化得到一个类别原型向量。查询图像素的分类不再依赖学习到的决策边界而是直接与原型向量算相似度。这种结构在竞赛中非常有效因为它的决策过程完全由支持集驱动天然适配小样本。我的具体选型是 ResNet-50 作为骨干网络去掉最后的全连接层保留到 conv5 输出的特征图。选择 ResNet-50 而非更深的 ResNet-101 的原因在于小样本场景下深网络更容易过拟合且梯度回传路径太长会导致原型提取时特征相关性下降。这里列一个选型对比表方案参数量小样本鲁棒性训练成本竞赛适用性Mask R-CNN 直接微调约 63M差RPN 不稳定高需要多阶段训练低U-Net 从头训练约 7M差容易过拟合中低ResNet-50 原型对齐约 23M好决策依赖支持集低高ResNet-101 原型对齐约 42M中特征过抽象中中方案确定后整个训练分为两个阶段第一阶段用基类数据集训练特征提取器损失函数是标准的语义分割交叉熵第二阶段冻结骨干网络用支持集样本训练原型提取模块与相似度度量模块。第二阶段只更新少量参数所以学习率要压得很低。2.3 数据划分基类预训练集、支持集、查询集怎么分数据划分是竞赛中拉开差距的第一道关卡也是后面所有技巧能够成立的前提。北极熊样本通常以视频帧序列的形式出现我建议按以下方式组织数据集第一步把全部北极熊影像按拍摄时间与地点分组。同一段连续视频中相邻几帧的北极熊姿态与背景变化极小这类强相关帧必须全部划入同一个小数据桶不能跨桶分配否则支持集和查询集之间会出现“伪相关”。第二步将数据桶按 6:1:3 的比例切分为基类微调桶、支持桶、查询桶。基类微调桶用于第二阶段的辅助训练支持桶用于生成原型查询桶用于最终评估。第三步在支持桶内做 K-shot 采样。如果支持桶有 120 帧K5则每轮训练随机取 5 帧作为支持集。每轮迭代都重新采样相当于在几十帧的基础上又做了一次数据增广。对于基类数据集常见做法是直接使用 COCO 数据集中包含“狗、猫、牛、羊”等家畜类别的图像这些类别与北极熊同属哺乳动物特征分布接近。用这些数据预训练第一阶段好处在于无需额外采集北极熊数据训练效率高。基类数据集每类取 20 到 50 张就够这是预训练不是主训练权重优先级远低于第二轮的原型学习。3. 用 PyTorch 跑通北极熊小样本分割三块核心代码与参数设置3.1 数据加载与掩膜预处理代码实现的第一个关键是数据加载北极熊图像的掩膜是单通道 PNG像素值 0 为背景1 为熊本体。读取后要统一缩放并与原图像做像素级对齐不能直接使用 PIL 的 resize因为掩膜的插值方式必须保持最近邻否则掩膜边缘会出现灰色过渡像素。下面是数据加载的核心代码import torch from torch.utils.data import Dataset from PIL import Image import numpy as np import torchvision.transforms as T class PolarBearDataset(Dataset): def __init__(self, image_list, mask_list, image_size256): self.image_list image_list self.mask_list mask_list self.image_size image_size def __len__(self): return len(self.image_list) def __getitem__(self, idx): # 读取原图转 RGB image Image.open(self.image_list[idx]).convert(RGB) # 读取掩膜modeL 保证是灰度图 mask Image.open(self.mask_list[idx]).convert(L) # 统一缩放到 256x256 image image.resize((self.image_size, self.image_size), Image.BILINEAR) # 掩膜必须用最近邻插值防止边缘出现中间值 mask mask.resize((self.image_size, self.image_size), Image.NEAREST) # 像素值归一化处理 mask np.array(mask) mask (mask 127).astype(np.float32) mask torch.from_numpy(mask).unsqueeze(0) # 图像归一化 image T.ToTensor()(image) return image, mask这段代码中掩膜缩放使用Image.NEAREST是关键细节。如果误用BILINEAR或BICUBIC掩膜边界会产生 0 到 1 之间的插值像素这些像素在后续计算损失时会被错误分类。图像归一化部分用的T.ToTensor()会把像素值从 [0, 255] 映射到 [0, 1]与预训练模型的输入规范保持一致。实际训练中我给每一轮迭代的查询集图像做了随机水平翻转与随机裁剪概率设为 0.5裁剪尺寸为 224×224 并保持宽高比不变。支持集图像不做增广因为原型是从精确掩膜中提取的几何扰动会让原型的位置信息失真。3.2 原型提取把支持集标注压成前景/背景类原型原型提取模块是小样本分割识别系统的核心它的任务是把支持集图像中掩膜区域的 CNN 特征聚合为一个向量。这里采用掩膜平均池化Masked Average Pooling前景原型与背景原型分开计算。背景并不直接取整张图特征的平均而是将掩膜外的像素特征做平均这样能让背景原型具备当前场景的空间全局性。import torch import torch.nn as nn import torch.nn.functional as F class PrototypeExtractor(nn.Module): def __init__(self): super().__init__() def forward(self, features, mask, eps1e-6): # features: [B, C, H, W] 查询集特征 # mask: [B, 1, H, W] 支持集掩膜0/1 值 mask F.interpolate(mask, sizefeatures.shape[-2:], modenearest) # 前景掩膜池化所有被掩膜覆盖的像素特征求平均 fore_mask mask fore_feats (features * fore_mask).sum(dim(2, 3)) fore_cnt fore_mask.sum(dim(2, 3)).clamp(mineps) fore_proto fore_feats / fore_cnt # 背景掩膜池化掩膜取反后求平均 back_mask 1.0 - mask back_feats (features * back_mask).sum(dim(2, 3)) back_cnt back_mask.sum(dim(2, 3)).clamp(mineps) back_proto back_feats / back_cnt # 返回 [B, C] 的前景原型与背景原型 return fore_proto, back_proto这段代码中的eps1e-6是为了防止某个原型区域的特征计数为零导致除零错误这在背景占比极小的裁剪图中经常发生。前景与背景原型最终会拼接成一个 [1, 2C] 的向量组作为后续像素相似度计算的对齐基准。参数设置上支持集一次前向只过一张图batch size 固定为 1这是出于原型稳健性的考虑。如果支持集同时过多张图并直接对特征取均值整只熊的姿态多样性会淹没在特征洪流中原型的判别力下降。正确做法是逐帧计算原型后对多个原型向量取平均而不是在特征图上直接取平均。3.3 损失函数与训练循环交叉熵 原型对比损失训练循环里第一阶段的基类预训练可以使用标准的多类交叉熵。第二阶段则要结合原型对比损失来拉近查询集特征与前景原型的距离拉开与背景原型的距离。具体实现时把查询集特征图上的每一像素当作一个独立样本计算它与前景原型和背景原型的余弦相似度再通过 Softmax 得到前景概率图。class FewShotSegmenter(nn.Module): def __init__(self, backbone, proto_extractor): super().__init__() self.backbone backbone self.proto_extractor proto_extractor def forward(self, support_img, support_mask, query_img): # 支持集特征与查询集特征共享同一个特征提取器 sup_feat self.backbone(support_img) que_feat self.backbone(query_img) fore_proto, back_proto self.proto_extractor(sup_feat, support_mask) # 计算查询特征和两个原型的余弦相似度 C fore_proto.shape[1] H, W que_feat.shape[-2:] que_flat que_feat.permute(0, 2, 3, 1).reshape(-1, C) sim_fore F.cosine_similarity(que_flat, fore_proto.unsqueeze(0), dim1) sim_back F.cosine_similarity(que_flat, back_proto.unsqueeze(0), dim1) # 堆叠成二分类 logits logits torch.stack([sim_back, sim_fore], dim1) logits logits.view(-1, 2, H, W) return logits def compute_loss(self, logits, query_mask): # 交叉熵损失query_mask 是 [B, 1, H, W] loss F.cross_entropy(logits, query_mask.squeeze(1).long()) return loss这里把背景相似度放在类别 0、前景相似度放在类别 1F.cross_entropy的输入不再需要额外做 Softmax因为它内部已经包含了 LogSoftmax。训练阶段我用的优化器是 Adam初始学习率 5e-4余弦退火调度总迭代轮数为 200 到 400 轮。基类预训练阶段的学习率是 1e-3骨干网络权重在第二阶段完全冻结只更新少数的 1×1 卷积和相似度度量的缩放因子。管线搭完之后一个值得注意的参数是特征图的输出尺寸。ResNet-50 输出下采样 32 倍的特征图也就是说输入 256×256 的图像特征是 8×8 的分辨率这对小目标而言太粗糙了。所以我额外接了一层简单特征金字塔提取 ResNet 的 conv2、conv3、conv4 三层特征分别计算原型再在 8 倍和 4 倍分辨率上融合预测。这个改动让分割识别系统在中等尺寸北极熊目标上的 mIoU 提升了约 7 个百分点。4. 推理流程与后处理从掩膜到连通的北极熊轮廓4.1 推理时的支持集选择策略训练完成后推理阶段仍然需要支持集的参与这与普通分割模型不同。竞赛中的常见做法是在每一段视频中随机选取一帧作为支持帧手动标注这一帧的北极熊掩膜然后对整段视频逐帧推理。支持帧的选择策略直接影响效果一般有静态支持与动态支持两种。静态支持是指固定一个支持帧全程使用适合目标外观变化不大的短序列动态支持则是每推理 20 帧重新从缓冲队列中选一帧更新原型适合熊转身、游泳等姿态剧烈变化的场景。我用过一个折中方案每 40 帧选一次支持帧选帧标准是当前帧预测掩膜的置信度均值置信度越接近 0.5 的帧越可能包含新的外观信息用它们更新原型池。这种做法能覆盖北极熊进入水中或爬上冰块时毛色反射率突然变化的情况。推理代码中支持帧的掩膜也需要缩放至与特征图相同的尺寸否则原型提取阶段会报维度不匹配的错误。这里有一个隐蔽问题模型越深特征图分辨率越小掩膜缩放比例越夸张ResNet-50 的 conv4 输出是输入尺寸的 1/32当原图是 1024×768 时掩膜要缩到 32×24此时误缩一位像素都会让原型错位。我的做法是先把原图与掩膜同时缩放到模型输入尺寸例如 384×288再做一次掩膜缩放两步比一步更稳。4.2 用条件随机场细化掩膜边界模型的直接输出是粗糙的像素概率图边缘部分常常出现毛刺小样本情境下尤其严重因为支持集没有足够多的边缘纹理信息让网络学会精确勾边。这类问题通常交给条件随机场CRF后处理解决。CRF 是一种图模型它以模型输出的概率分布为一元势能以像素颜色与位置的相似度为二元势能通过迭代推理让相邻像素的标签趋于一致。竞赛中我用的是 DenseCRF它可以在几毫秒内完成一张图的迭代。实现时我借助了 PyTorch 自带的计算图自己写了简化版本核心参数只有两个空间核权重w23与颜色核权重w110。这两个参数控制着 CRF 对颜色纹理的信任程度。北极熊的毛色如果与背景雪地接近颜色核权重过大会导致整只熊被误删为背景权重过小则边界细化效果为零。常用调试方式是先固定空间核权重为 3再递增颜色核权重观察熊的背脊线是否被保留。另外提醒一点CRF 不是万能的它只能修复小尺度锯齿如果原型本身嵌入错位导致整片区域分类错误CRF 救不回来。4.3 性能评估mIoU 与像素级 F1 的计算口径竞赛评测标准通常以 IoU 为核心但只报 mIoU 会掩盖问题。小样本分割识别系统的性能波动很大支持集与查询集的分布差异可以直接导致的 mIoU 从 0.65 掉到 0.45所以评估必须有分层结构。第一层是全局 mIoU所有测试图像的混淆矩阵累加后计算 IoU。第二层是分桶 mIoU按北极熊在画面中的面积比例分成大目标16%、中目标4%-16%、小目标4%三档分别统计。第三层是交叉验证结果每次随机采 5 组支持集报告 mIoU 的均值与标准差。我自己的项目经验是小样本模型常出现“大目标高分、小目标低分”的严重断层原因在于小目标在特征图上只有不到 10 个有效像素原型提取时池化到了大量背景信息。因此我在评估时格外关注小目标档的 mIoU并把这个指标与整体 mIoU 一起呈现在竞赛答辩材料中这也容易得到评委认可。以下表格是一个典型的小样本分割识别系统评估样例用来帮助大家理解不同 K-shot 值下性能的下降趋势K-shot大目标 mIoU中目标 mIoU小目标 mIoU全图 mIoUK10.620.470.220.51K30.680.550.300.58K50.710.590.340.62看这张表就能直观感受到小样本学习的边界效果增加 K 值对中等目标提升明显对极小目标则收效甚微。如果比赛报告只写一个全图 mIoU这种分层信息就完全被掩盖了。5. 北极熊分割识别系统最容易翻车的 4 个坑现象、原因与解法5.1 验证集 mIoU 虚高支持集与查询集来自同一段视频现象模型在本地验证集上 mIoU 高达 0.78但主观目视发现很多帧分割结果与真实情况差异很大明显的数据与现象不一致。原因北极熊视频连续帧之间目标位移小背景几乎不变特征图之间的余弦相似度天然高。当支持集和查询集取自同一段连续视频时模型本质上是在“回忆”前一帧而不是学习判别特征这就是数据泄漏。解决按数据桶划分数据集同一段视频内帧序列只能进入同一个桶桶级别划分比例 6:1:3然后确保评估时使用的支持集与查询集来自不同桶。这条规则应当在代码里强制校验而不是靠口头约定。5.2 背景类主导损失雪地与熊的像素分布严重失衡现象训练时损失函数曲线平滑下降但北极熊区域的召回率极低漏检率高达 40% 以上几乎处于全检背景的状态。原因自然场景中北极熊目标占画面面积通常不到 10%背景类像素占绝对优势。标准交叉熵损失会被背景类主导模型学到“全部预测为背景”即可获得低损失。解决在损失函数中给前景类添加权重因子常见设置是class_weight [1.0, 5.0]或者使用 Dice Loss 与交叉熵按 0.5:0.5 比例混合。我实验发现对小目标尤其有效的做法是在计算损失前把前景区域的损失做单独的上采样放大等效于给难例更高权重。5.3 冻结骨干网络后支持集损失不下降学习率没按层调整现象第二阶段训练开始后支持集损失在 50 轮内纹丝不动余弦相似度输出始终集中在一个固定值附近。原因骨干网络已经冻结只有最后的少量卷基层参与训练但学习率仍然沿用第一阶段 1e-3 的数值。新初始化的 1×1 卷积层在这么高的学习率下梯度振荡剧烈无法收敛到有效解。解决第二阶段改用分组学习率骨干网络梯度不计算requires_gradFalse新加的解码层与原型提取层使用 5e-4 学习率。同时加入梯度裁剪clip_grad_norm_(max_norm0.5)防止相似度模块初始阶段的梯度爆炸。5.4 直接取 argmax 导致目标碎片化缺少空间连续性约束现象模型预测的掩膜呈现“斑驳”状态北极熊身体被分割成几十个碎块按压为连通域后发现最大的连通域只占全部预测像素的 30%。原因像素级独立的 Softmax 决策没考虑空间相关性。小样本提取的特征本身具备一定噪声加上原型匹配时背景特征与前景特征在颜色相近区域区分度低逐像素解码自然产生碎片。解决预测输出后先做一次形态学闭运算核大小设为 3×3然后保留置信度大于 0.65 的连通域并剔除面积小于 500 像素的孤立区域。经过这套后处理连通域占比从 30% 提升到 80% 以上mIoU 在测试集上可提升 4 到 6 个百分点。6. 再进一步用小样本分割技巧拉高北极熊系统的落地上限6.1 余弦相似度加温度缩放让置信度分布更可信原型匹配输出的余弦相似度范围在 [-1, 1] 之间直接做 Softmax 会导致绝大多数像素的置信度集中在 0.5 附近不利于阈值分割。常见做法是在相似度上乘以一个温度系数 τ经验值取 5.0。温度系数实际上在做最邻近类别匹配时放大类间差异让小目标区域也能形成峰值响应。代码上只需要把sim_fore sim_fore * 5.0即可不需要改动网络结构。6.2 推理时动态扩增支持集查询帧回注推理时发现当前帧预测的置信度超过 0.9就可以把这个帧的预测掩膜作为伪标注加入支持集缓冲池供后续帧重新生成原型。这种做法本质上是自举式的半监督学习能在视频序列的持续推理中逐步提升原型质量。注意伪标注必须经过上述连通域筛选最终只保留面积最大、置信度最高的那一块避免错误掩膜污染原型池。6.3 交叉验证评估小样本性能报告均值和方差小样本分割识别系统的单次评估结果随机性大这是小样本学习的典型特性不是模型缺陷。因此在最终汇报时固定 K-shot在 5 组不同的支持集随机抽样下各跑一轮评测报告 mIoU 的均值与标准差。只报最好一次的 mIoU 没有参考价值评委与工程团队真正想知道的是模型的稳定性边界。这个习惯也延续到了我后续的工业项目中每次小样本模型验收我都会额外跑一次交叉验证并给出置信区间很多模型都是在稳定性和边界测试中暴露出来的真实问题。在竞赛答辩和项目复盘中我常常会保留当前最优的一组支持集特征可视化图。看着那些经过掩膜池化得到的前景原型在不同光照条件下如何偏移你就能真切理解小样本学习的局限与空间它不是一个黑匣子而是一个受数据分布约束的匹配过程。希望这套从方案设计到坑点修复的路径能帮到你。本文还有配套的精品资源点击获取
返回列表