滑动窗口内存优化与流式落盘

滑动窗口内存优化与流式落盘

学习目标

读完本章后，你将能够：

• 理解长文档处理中的内存瓶颈问题
• 掌握滑动窗口分页处理的实现原理
• 分析流式落盘如何降低内存峰值
• 设计适合超长文档的处理方案

前置知识

• 无

---

1. 长文档处理的内存挑战

处理上万页的 PDF 文档时，主要内存消耗来自：

消耗项	说明	单页估算
页面图片	PDF 页面渲染为图片（300 DPI）	10-20 MB
模型输入	图片送入模型推理	5-10 MB
中间结果	布局检测、OCR、MFR、TSR 结果	2-5 MB
最终输出	Markdown/JSON	0.1 MB

10,000 页文档的理论内存：10,000 × 20 MB = 200 GB → 不可接受

2. 滑动窗口机制

2.1 核心思路

不一次性加载所有页面，而是维护一个固定大小的处理窗口：

2.2 实现伪代码

def process_with_sliding_window(
    pdf_doc: PdfDocument,
    window_size: int = 64,
):
    total_pages = len(pdf_doc)
    results = []

    for start_idx in range(0, total_pages, window_size):
        end_idx = min(start_idx + window_size, total_pages)

        # 1. 仅加载窗口内的页面
        window_images = load_images_from_pdf_doc(
            pdf_doc,
            start_page_id=start_idx,
            end_page_id=end_idx,
        )

        # 2. 处理窗口内页面
        window_results = analyze_pages(window_images)

        # 3. 流式落盘：将已完成的结果写入磁盘
        for result in window_results:
            write_to_disk(result)

        # 4. 释放窗口内存
        del window_images
        clean_memory()  # 清理 GPU/CPU 缓存

        # 5. 合并结果
        results.extend(window_results)

    return results

关键参数：

• window_size：默认 64 页，可通过配置调整
• clean_memory()：清理 torch.cuda.empty_cache() 和 gc.collect()

2.3 跨页元素处理

滑动窗口的挑战是跨页元素（如跨页表格、连续段落）：

# 窗口边界处的跨页检测
def handle_cross_page_elements(current_window, next_window):
    # 检查当前窗口最后一页是否有未完成元素
    last_page = current_window[-1]

    for block in last_page.blocks:
        if block.type == "table" and block.is_truncated:
            # 在下一个窗口中寻找表格续页
            continuation = find_table_continuation(next_window[0], block)
            if continuation:
                merge_tables(block, continuation)

3. 流式落盘

3.1 原理

传统方案：所有页面处理完后一次性写入

页面1 → 页面2 → ... → 页面N → 一次性写入磁盘
                                 ↑
                            全部结果在内存中

流式落盘方案：每个窗口处理完立即写入

窗口1处理完 → 立即写入磁盘 → 释放内存
窗口2处理完 → 立即写入磁盘 → 释放内存
窗口3处理完 → 立即写入磁盘 → 释放内存

3.2 实现

def streaming_write(results, output_dir):
    """流式落盘：每处理完一个窗口就写入"""
    for page_result in results:
        # 写入中间 JSON
        append_to_middle_json(page_result, output_dir / "middle.json")

        # 写入图片
        for img in page_result.images:
            img.save(output_dir / "images" / img.name)

        # 可选：增量生成 Markdown
        append_to_markdown(page_result, output_dir / "output.md")

关键设计：

• 中间 JSON 使用追加写入（append mode）
• 图片文件即时写入磁盘
• 最终 Markdown 可以增量生成

4. Batch 推理优化

4.1 Pipeline 的 Batch 推理

Pipeline 引擎支持对窗口内页面进行批量推理：

def batch_analyze(window_images, batch_size=8):
    """批量推理 + 流式落盘"""
    results = []

    for i in range(0, len(window_images), batch_size):
        batch = window_images[i:i+batch_size]

        # 批量推理
        batch_results = model.predict_batch(batch)

        # 流式落盘
        for result in batch_results:
            write_to_disk(result)
            results.append(result)

        # 清理批次内存
        clean_memory()

    return results

4.2 VLM 的 Batch 推理

VLM 引擎根据后端不同支持不同的 batch 策略：

后端	Batch 支持	说明
transformers	可配置	batch_size 参数控制
vllm	自动	vLLM 自动批处理
lmdeploy	自动	LMDeploy 自动批处理
mlx	有限	macOS 内存限制

5. 线程安全推理

v3.0 版本引入了线程安全优化：

# 使用 threading.RLock 保护模型推理
class ModelSingleton:
    _lock = threading.RLock()

    def predict(self, image):
        with self._lock:
            return self.model.predict(image)

收益：

• 支持多线程并发推理
• 配合 mineru-router 实现多卡部署
• 提升单机吞吐量

6. 陷阱与对策

陷阱	对策
窗口大小设置过大导致 OOM	根据可用显存调整 processing_window_size
跨页表格合并失败	确保表格续页检测功能正常
流式落盘时磁盘 I/O 瓶颈	使用 SSD，或增加落盘缓冲
Batch 大小与显存不匹配	减小 batch_size，或切换到纯 CPU 模式
内存泄漏导致逐页增长	确保每批次后调用 clean_memory()

7. 设计权衡

窗口大小选择

窗口大小	优势	劣势
小（16页）	内存占用低	批处理效率低，频繁 I/O
中（64页）	平衡点	需要约 1-2 GB 额外内存
大（256页）	批处理效率高	内存占用高

取舍：默认 64 页，适用于大多数场景。大内存机器可增加到 128 页提升速度，低内存机器可减少到 16 页。

---

参考来源

• MinerU 源码：mineru/backend/pipeline/pipeline_analyze.py
• MinerU CHANGELOG v3.0：滑动窗口优化与流式落盘