20260317 update

我发现了一个更小巧的公式识别模型：Texo。因此当前仓库仅作为学习使用了。

√ TrOCR Formula Recognition

❓缘由：看到 UniMERNet 的工作，从他们发布的模型存储大小（4.91G）来看，实在太重了。同时，他们也发布了一个很大很全的公式识别数据集：UniMER_Dataset。

🎯 于是，想着基于 TrOCR + UniMER-1M 数据集，训练一个小而美的公式识别数据集。

仓库将 UniMERNet 作为 Baseline，目标是超过 UniMERNet，同时模型要小很多。

仓库 dataset 目录下为 UniMER-1M 的 Tiny 版，只用来测试程序使用。

⚠️注意事项

• 使用 transformers 训练前，需要在 import torch 前，指定 CUDA_VISIBLE_DEVICES，否则会卡住。
• 以下实验数据，除 Exp1_1 外，其他的暂时都没有添加 HME100K 数据集
• 所有实验均采用 microsoft/trocr-small-stage1 作为预训练模型训练的。

TODO

• 给出速度基准
• 推理采用 Flash Attention 加速。（transformers==4.44.2 中 VisionEncoderDecoderModel 不支持）
• 转 ONNX 模型，并比较推理速度
• 尝试使用 xformers 来优化推理速度

🔬 实验记录

实验表格来自 UniMERNet Table 5

Method	SPE-BLEU↑	SPE-EditDis↓	CPE-BLEU↑	CPE-EditDis↓	SCE-BLEU↑	SCE-EditDis↓	HWE-BLEU↑	HWE-EditDis↓
Pix2tex	0.873	0.088	0.655	0.408	0.092	0.817	0.012	0.920
Texify	0.906	0.061	0.690	0.230	0.420	0.390	0.341	0.522
UniMERNet	0.917	0.058	0.916	0.060	0.616	0.229	0.921	0.055
Exp1	0.815	0.121	0.677	0.259	0.589	0.227	0.150	0.520
Exp1_1	0.883	0.07	0.810	0.122	0.489	0.262	0.900	0.06
Exp2	0.798	0.132	0.677	0.259	0.589	0.227	0.150	0.520
Exp3	0.813	0.127	0.682	0.263	0.302	0.231	0.166	0.540
Exp4	0.873	0.077	0.801	0.130	0.550	0.238	0.092	0.469
Exp5	0.846	0.201	0.823	0.134	0.418	0.553	0.05	0.6724
Exp5_1	0.819	0.119	0.682	0.249	0.595	0.230	0.179	0.512
Exp6	0.812	0.116	0.676	0.253	0.657	0.210	0.342	0.404
Exp7	0.817	0.117	0.679	0.251	0.817	0.117	0.781	0.148
Exp8	0.886	0.07	0.822	0.108	0.633	0.217	0.897	0.07
Exp9	0.862	0.10	0.740	0.180	0.639	0.211	0.826	0.119
Exp10	0.900	0.07	0.841	0.09	0.594	0.227	0.912	0.06

Exp	说明
Exp1	首次基于 UniMER-1M 训练，采用预训练模型是 microsoft/trocr-small-stage1 采用 TrOCR 默认 Tokenizer
Exp1_1	基于 Exp1，控制单一变量：训练 30 个 Epoch by limaopeng1
Exp2	更改 LaTex-OCR 方法用的 BPE Tokenizer
Exp3	修复 Exp2 中 model 配置 bug
Exp4	与 Exp3 相比，单一变量：epoch=1 → epoch=5
Exp5	与 Exp1 相比，单一变量：epoch=1 → epoch=10
Exp5_1	补充实验，修复 Exp5 中，去掉 text 前后加了 BOS 和 EOS 的地方，只跑一个 epoch
Exp6	与 Exp5_1 相比，单一变量：参考 UniMERNet 源码，增加数据增强
Exp7	与 Exp6 相比，单一变量：增加 HME100k 数据集
Exp8	与 Exp7 相比，单一变量：epoch=1 → epoch=10
Exp9	与 Exp8 相比，单一变量：增加 fusion-image-to-latex-datasets 数据集（3069505）, Epoch=1
Exp10	与 Exp9 相比，单一变量：epoch=1 → epoch=10 (fusion-image-to-latex-dataset 3467214)

🦩 Checkpoint

• Exp5_1
• Exp8
• Exp10

🔢 Dataset

⚠️注意：仓库中 dataset 目录下为示例，完整数据集需自行下载补充。

UniMER_Dataset 完整的 UniMER 目录结构如下：

dataset
├── UniMER-1M
│   ├── images
│   └── train.txt
└── UniMER-Test
    ├── cpe
    ├── hwe
    ├── sce
    ├── spe
    ├── cpe.txt
    ├── hwe.txt
    ├── sce.txt
    └── spe.txt

训练集总共 1061,791 LaTeX-Image pairs。

测试集由 4 种类型公式组成，总共 23757 张图像：

• Simple Printed Expressions (SPE): 6,762 samples
• Complex Printed Expressions (CPE): 5,921 samples
• Screen Capture Expressions (SCE): 4,742 samples
• Handwritten Expressions (HWE): 6,332 samples

各个种类示例图像如下：

🔢 其他数据集和项目

• fusion-image-to-latex-datasets
• TexTeller

📞 微信交流群

微信关注公众号：RapidAI, 后台回复“公式识别”即可进群

📚 Reference

• UniMERNet
• TrOCR-Handwritten-Mathematical-Expression-Recognition
• Transformers-Tutorials