企业官网建设流程全解析

DeepSeek-OCR-2 Transformers推理教程：flash_attention_2 + bfloat16如何加速单卡OCR

【免费下载链接】DeepSeek-OCR-2Visual Causal Flow项目地址: https://gitcode.com/gh_mirrors/de/DeepSeek-OCR-2

本教程带你用Transformers在单张显卡上跑通DeepSeek-OCR-2的 OCR 推理，并通过flash_attention_2注意力后端 +bfloat16半精度两项关键配置实现显存减半、推理提速。DeepSeek-OCR-2 是 DeepSeek 推出的多模态 OCR 模型，主打"Visual Causal Flow"（视觉因果流），能把扫描文档转成结构化 Markdown，支持动态分辨率和自由 OCR 两种模式。

一、DeepSeek-OCR-2 是什么？

与传统 OCR 不同，DeepSeek-OCR-2 将语言模型（LM）当作视觉编码器：右侧的 DeepEncoder V2 使用 Qwen2 500M 作为"LM as Vision Encoder"，配合可学习的因果查询（learnable query）生成新的阅读顺序，比左侧基于 CLIP ViT 的初代 DeepEncoder 更贴近人类阅读逻辑。

更多技术细节可参考官方论文 DeepSeek_OCR2_paper.pdf 与 README.md。

二、一键安装：单卡推理环境搭建

官方环境为CUDA 11.8 + PyTorch 2.6.0，Python 3.12.9：

git clone https://gitcode.com/gh_mirrors/de/DeepSeek-OCR-2 conda create -n deepseek-ocr2 python=3.12.9 -y conda activate deepseek-ocr2

安装 PyTorch、模型依赖与 FlashAttention 核心组件：

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu118 pip install -r requirements.txt pip install flash-attn==2.7.3 --no-build-isolation

💡 其中flash-attn是flash_attention_2的底层依赖，缺少它会导致注意力后端报错。依赖清单见 requirements.txt。

三、Transformers 推理：核心 10 行代码

仓库提供了现成的 HF 推理脚本 run_dpsk_ocr2.py，核心逻辑如下：

from transformers import AutoModel, AutoTokenizer import torch model_name = 'deepseek-ai/DeepSeek-OCR-2' tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) model = AutoModel.from_pretrained( model_name, _attn_implementation='flash_attention_2', # ⚡ 注意力加速 trust_remote_code=True, use_safetensors=True ) model = model.eval().cuda().to(torch.bfloat16) # ⚡ 半精度加速 prompt = "<image>\n<|grounding|>Convert the document to markdown. " res = model.infer( tokenizer, prompt=prompt, image_file='your_image.jpg', output_path='your/output/dir', base_size=1024, image_size=768, crop_mode=True, save_results=True )

运行方式：

cd DeepSeek-OCR2-master/DeepSeek-OCR2-hf python run_dpsk_ocr2.py

两个加速参数在 run_dpsk_ocr2.py 中只占两行，却是单卡跑动的关键。

四、加速原理：两项配置为什么有效？

1️⃣ flash_attention_2：更快的注意力计算

默认注意力实现需要显式存下完整的 N×N 注意力矩阵，长序列下显存占用高、速度慢。_attn_implementation='flash_attention_2'启用 FlashAttention-2 内核：

• 分块计算（tiling）：注意力矩阵不落显存，显存占用从 O(N²) 降到 O(N)
• IO 感知设计：减少显存读写次数，单卡吞吐显著提升
• 动态分辨率友好：DeepSeek-OCR-2 一张图可产生数百至上千个视觉 token，序列越长，FlashAttention 收益越大

2️⃣ bfloat16：显存与计算双减半

.to(torch.bfloat16)将模型权重从 32 位浮点压缩为 16 位：

• 模型显存占用近乎减半，让单卡（如 16GB 显存）从容加载
• A100/RTX 30 系及以上 GPU 对 bf16 有硬件加速，矩阵运算速度翻倍
• bf16 动态范围与 fp32 相同，比 fp16 更不易溢出，推理结果稳定

两者叠加，典型场景下可节省约一半显存并获得明显的端到端提速。

五、Prompt 怎么选？两种模式对比

六、infer 关键参数速查

• base_size=1024/image_size=768：动态分辨率的两个基准尺寸。默认方案为 (0-6)×768×768 + 1×1024×1024，对应 (0-6)×144 + 256 个视觉 token，图文越长 token 越多，这正是启用flash_attention_2的原因
• crop_mode=True：启用区域裁切，适合版面复杂的长文档
• save_results=True：自动保存识别结果到output_path

七、常见问题 FAQ

Q：报flash_attn相关错误怎么办？确认已安装匹配版本：pip install flash-attn==2.7.3 --no-build-isolation，且 PyTorch 为 cu118 构建的 2.6.0。

Q：想同时跑 vLLM 和 Transformers 代码？可以共用一个环境，遇到vllm requires transformers>=4.51.1的告警可忽略（见 README.md Install 章节的说明）。vLLM 推理入口在 DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/run_dpsk_ocr2_image.py，配置项在 config.py。

Q：显存还是不够？优先确认已启用 bf16；必要时通过CUDA_VISIBLE_DEVICES指定空闲 GPU，并减小输入图片尺寸。

总结

DeepSeek-OCR-2 的 Transformers 推理只需三步：装好 CUDA 11.8 + flash-attn 环境 → 加载模型时指定flash_attention_2→ 转成bfloat16上卡。两项加速配置合计仅两行代码，却能显著降低显存占用并提升单卡 OCR 速度，是新手低成本跑通 DeepSeek-OCR-2 的完整路径。

【免费下载链接】DeepSeek-OCR-2Visual Causal Flow项目地址: https://gitcode.com/gh_mirrors/de/DeepSeek-OCR-2