第6名解决方案 - deplot与UNet后处理

感谢Kaggle和Benetech举办这场有趣的比赛。这是一个与@darraghdog共同努力的团队成果。

TLDR

我们的解决方案是对单个deplot模型进行4次不同种子融合的结果，适用于所有图表类型。对于散点图，我们通过U-Net点分割模型预测的散点数量进行额外的后处理。

数据预处理

Benetech数据集

训练时我们使用了全部Benetech提取的数据，以及约20%基于最困难样本选择的生成数据。这显著加快了训练流程。
真实数值的标签包含大量数字，这与基于pix2struct的模型（如deplot）的损失函数（交叉熵）不太匹配。因此，对这些数值进行舍入处理非常重要。为此，我们在刻度标签轴内将标签分桶至约5个区间。例如，下图中的分桶会是0, 40000, 80000, 120000, 160000, 200000, 240000… 等等，而y轴标签学习的目标则是840000, 960000, 760000, 700000, 580000, 460000, 等等。对于没有刻度标签的真实值标签，则根据数据系列的最大最小值范围进行类似的舍入处理，使用约40个分桶。

在pix2struct模型中，标签以键值对序列的形式学习，类似于deplot的训练方式。以上述例子为例，学习序列为：
<pad>Malaysia;960000|Malawi;840000|Maldives;760000|Mali;700000|Mauritania;660000|Melanesia;580000|Malta;580000|Mauritius;540000|Martinique;460000|Mayotte;420000<\\s>

我们还使用了一个独立的解码器来学习反向的键值对序列。该解码器与相同的编码器相连，并作为辅助损失加入，使模型不会过度关注图像的任一侧。由于deplot的主要计算开销在编码器部分，因此双解码器并未带来显著的性能负担。

外部数据

本次比赛中，我们使用了从PubMed文章中商业可用的子集来训练模型。发现ICDAR数据集中约80%的标注样本来自该商业数据集，并确保其使用符合相关许可条款。为扩充训练数据，我们采用伪标签技术后进行人工审核。具体来说，我们对约8,000个来自PubMed文章的条形图和折线图样本进行了伪标签标注和人工审核。伪标签在条形图上的结果相当准确，只需验证图表类型是否正确（如移除箱线图）。折线图则需要一些修正。我们还使用U-Net模型从PubMed伪标注了约1,000个散点图样本，然后手动补充缺失的数据点。使用VGG图像标注工具，这一人工标注过程出人意料地高效快捷。

数据生成

我们采用@brendanartley优秀的公开内核，主要通过两种方式：
首先，我们使用该内核可下载的50万张图像对deplot模型进行预训练，相比直接使用Hugging Face上的原始检查点获得了小幅提升。

此外，我们对内核进行了一些调整，生成更复杂的散点图（更多点、不同标记、更多重叠等），并输出每个散点的x/y坐标。我们用该脚本生成了5万张散点图，并额外用于训练分割模型。

模型

使用deplot预测条形图、折线图和散点图；对于散点图，我们还额外使用了U-Net模型。

Deplot训练了20个epoch，学习率约为5e-5，批大小为8。发现Pixeldropout数据增强能有效防止过拟合，未使用其他图像增强方法。实验阶段使用1024个patch，训练时间缩短了2倍，仅轻微影响分数。损失函数结合了键值序列的交叉熵和图表类型预测。在deplot的第一个token之上添加线性层，以辅助方式预测图表类型。

对于散点图，我们采用U-Net（EfficientNet-B7编码器）+非极大值抑制（NMS）来预测散点数量，并在deplot预测点数不准时执行简单而有效的后处理：

• 如果deplot预测的点数过多，则舍弃最后几个点
• 如果deplot预测的点数过少，则用均值填充缺失值

下图展示了U-Net训练的分割掩码。使用散点标注，我们创建具有10像素高斯半径的"平滑"点并绘制到空白图像上。

模型集成

对于deplot，我们在每个贪婪解码步骤中平均化解码器预测结果进行内部集成。对于Unet，我们简单地在执行NMS前平均预测掩码。

消融研究（Public LB分数）

下表概述了各个组件对解决方案的贡献：

使用的工具/代码库

• Unet: Segmentation models pytorch + timm编码器
• Deplot: huggingface
• 日志/可视化: neptune.ai
• 数据增强: albumentations
• 标注工具: VGG Image Annotator工具

感谢阅读，欢迎提问！