工具-MinerU高显存占用解决方案

在文档智能处理（Document AI）领域，MinerU 是一款强大的开源 PDF 解析工具，能够精准提取文档中的 Markdown 内容、公式和表格。然而，在实际工程落地中，直接将几百页的 PDF 扔给模型处理，往往会面临一个棘手的问题——显存（VRAM）爆炸。

本文将结合一段 Python 实战代码，深入剖析为什么大文件会导致显存溢出，并阐述一种基于“单页拆分”的优化策略及其背后的原理。

1. 痛点：为什么大文件会让显存“烟花”？

MinerU（以及大多数基于 Transformer 或视觉编码器的模型）在处理文档时，需要将输入内容转换为高维向量表示。要回答“为什么拆分后显存就不炸了”，我们需要深入到深度学习模型的显存占用机制。

1.1 显存都去哪儿了？

在深度学习推理（Inference）过程中，显存主要被以下几部分占用：

模型权重（Model Weights）：
这是固定的。不管你输入是一个字还是一万个字，加载模型本身（如 InternVL、LayoutLM 等）就需要占用几个 GB 的显存。这部分是“入场费”。
中间激活值（Activations）与临时缓冲区：
这是显存爆炸的罪魁祸首。
- Transformer 的注意力机制（Self-Attention）：这是 Transformer 的核心，但它的计算复杂度是 **$O(L^2)$**，其中 $L$ 是输入序列的长度。
  - 如果你输入一页 PDF，提取出 1,000 个 token（或图像 patch），计算量约为 $1,000^2 = 1,000,000$。
  - 如果你一次性输入 100 页 PDF，序列长度变为 100,000，计算量约为 $100,000^2 = 10,000,000,000$。
  - 注意：长度扩大 100 倍，计算量和显存需求（存储 Attention Matrix）会扩大 10,000 倍！
- 中间特征图（Feature Maps）：对于视觉模型（Vision Encoder），输入图像越大（分辨率越高），每一层卷积或 Transformer Block 生成的特征图就越大。大文件往往意味着超长、超大的图片输入，每一层的中间结果都需要显存来存放。

1.2 为什么“拆分”能解决问题？

当你把一个 100 页的 PDF 拆分成 100 个 1 页的 PDF 时：

序列长度（L）被截断：每次推理，模型只需要处理 1 页的内容。$L$ 始终保持在一个很小的数值（比如 1,000）。
显存占用重置：处理完第 1 页后，该次推理产生的中间激活值会被释放，显存被清空，用于处理第 2 页。
峰值显存锁定：整个任务期间，显存的峰值不再取决于总页数，而是取决于最复杂的那一页。

这就好比只有一辆卡车（显存），要运送 100 吨货物（100页）。

不拆分：试图把 100 吨货一次性装上车 —— 车压塌了（OOM）。
拆分：每次只装 1 吨，跑 100 趟 —— 轻松完成。

1.3 灵魂拷问：如果我并发处理小文件，显存会炸吗？

这是一个非常好的底层问题。答案是：会，而且可能会炸得更惨。

如果你使用多线程/多进程，同时向 GPU 发送 10 个单页 PDF 的请求：

Batch Size 效应：在 GPU 看来，这等同于你把 Batch Size 从 1 变成了 10。
显存叠加：虽然每个任务只占 1 份激活值显存，但 10 个任务并行，就需要 10 份显存。
- 公式大致为：Total_VRAM = Model_Weights + (Single_Page_Activations * Concurrency_Count)
结论：拆分后的并发数是受限于显存大小的。你不能无限并发。通常对于大模型，单卡并发数（Batch Size）设为 1 是最稳妥的，或者需要经过严格的压力测试来确定最大安全并发数。

2. 解决方案：化整为零，逐页击破

既然“一口吃成个胖子”行不通，最好的办法就是“细嚼慢咽”。我们采用 按页拆分（Page-by-Page Splitting） 的策略：

拆分：利用 PDF 处理库（如 pypdf）将原始的大文件 PDF 拆解为一个个独立的单页 PDF。
处理：每次只将一个单页 PDF 发送给 MinerU 服务进行解析。
聚合：将解析回来的 Markdown 结果按顺序拼接。

这种方式将显存的复杂度从 **O(N)**（N为总页数）降低到了 **O(1)**（仅与单页复杂度相关）。

3. 代码核心逻辑解析

让我们深入 utils/pdf_cntent_mineru.py 中的 ocr_pdf_per_page 函数，看看这一策略是如何落地的。

3.1 核心流程

def ocr_pdf_per_page(pdf_path: str | pathlib.Path) -> str:
    results: List[str] = []
    # ... (省略部分逻辑)
    with open(pdf_path, "rb") as infile:
        reader = pypdf.PdfReader(infile)
        total_pages = len(reader.pages)
        
        # 循环遍历每一页
        for page_index in range(total_pages):
            # ... (省略停止逻辑)
            
            # 【关键步骤】创建临时文件，只包含当前页
            with tempfile.NamedTemporaryFile(suffix=".pdf", delete=True) as tmp:
                writer = pypdf.PdfWriter()
                writer.add_page(reader.pages[page_index]) # 提取单页
                writer.write(tmp)                         # 写入临时文件
                tmp.flush()                               # 确保数据落盘
                
                # 调用 MinerU 接口处理这个单页文件
                res = get_ocr(tmp.name, page_index)
                
                # 获取解析结果
                res_str = res.get('results', {}).get(
                    str(page_index), {}).get('md_content', '')
                
                results.append(res_str)
                
    return ''.join(results)

3.2 关键技术点

临时文件（Tempfile）妙用：
代码使用了 tempfile.NamedTemporaryFile。这非常巧妙，因为它创建了一个系统级的临时文件，像普通文件一样拥有路径（tmp.name），可以被 requests 库读取并上传。且 delete=True 保证了上下文退出后文件自动删除，不会产生垃圾文件。
流式内存控制：
pypdf 的 PdfReader 是懒加载的，它不会一次性把整个 PDF 加载到内存。在循环中，我们每次只提取 reader.pages[page_index] 对象，并写入一个新的微型 PDF。这意味着无论源文件有 1GB 还是 10MB，Python 端的内存占用都极低。
业务逻辑融合（Early Stopping）：
代码中还包含了一个有趣的逻辑：
1
2
if _contains_amount_info(res_str):
stop_after = 5
这是基于页维度的流式处理带来的额外红利。如果我们在第 3 页就找到了需要的“金额信息”，我们可以设置一个缓冲（再读5页）然后直接 break。如果是全量上传，必须等所有页解析完才能进行业务判断，浪费了大量的 GPU 算力。

4. 算法优势总结

采用这种“单页拆分”算法，对比全量上传，拥有显著优势：

维度	全量上传 (Original)	单页拆分 (Current Strategy)
显存占用	极高且不可控。随页数增加线性或指数增长。	低且恒定。仅取决于最复杂的那一页，显存峰值可控。
稳定性	大文件极易导致 OOM 崩溃，服务不可用。	极其稳定，即使处理千页文档也能稳如泰山。
容错性	整个任务要么成功，要么失败。	粒度细化到页。若第 50 页解析失败，不影响前 49 页的结果。
业务灵活性	必须等待全量解析完成。	支持流式处理和提前终止，大幅节省时间和算力。
并发扩展性	难以拆分任务。	天然支持并发。未来可以改造为多线程同时发送多个单页请求，利用集群算力。