本模型是基于 ConvNeXt-Tiny + ViTSTR (Vision Transformer for Scene Text Recognition) 架构的手写体 OCR 文字识别模型,由 DAMO Academy (阿里巴巴达摩院) 提供。模型采用 ConvNeXt 作为视觉骨干网络提取图像特征,结合 Vision Transformer (ViT) 进行序列建模,最终通过全连接层分类输出识别结果。本模型专门针对手写体文字进行了优化训练。
模型由三部分组成:
本模型已成功适配华为昇腾 Ascend NPU,支持在 NPU 上进行推理。主要适配工作:
# 使用清华 PyPI 镜像
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision Pillow numpy
# 如需 NPU 支持,请安装 torch_npu(根据你的 CANN 版本选择对应包)模型通过 ModelScope 自动下载缓存到 ~/.cache/modelscope/hub/models/cycloneboy/cv_convnextTiny_ocr-recognition-handwritten_damo/,包含:
pytorch_model.bin: 模型权重 (73.3 MB)vocab.txt: 字符词汇表 (7643 个字符)configuration.json: 模型配置python3 inference.py --device cpu --image /path/to/your/image.pngpython3 inference.py --device npu --image /path/to/your/image.png| 参数 | 默认值 | 说明 |
|---|---|---|
--image | /opt/atomgit/test_ocr_image.png | 输入图像路径 |
--device | cpu | 运行设备 (cpu 或 npu) |
--model-dir | 缓存路径 | 模型文件目录 |
使用测试图像 test_ocr_image.png (399×78 像素) 进行推理:
============================================================
Model: cv_convnextTiny_ocr-recognition-handwritten_damo
Device: CPU
Input: /opt/atomgit/test_ocr_image.png
============================================================
Loading model...
Vocabulary size: 7643
Input tensor shape: torch.Size([1, 1, 32, 163])
------------------------------------------------------------
OCR RESULT:
------------------------------------------------------------
拍江买
------------------------------------------------------------
Inference time on CPU: 0.0750s
Done.============================================================
Model: cv_convnextTiny_ocr-recognition-handwritten_damo
Device: NPU
Input: /opt/atomgit/test_ocr_image.png
============================================================
Loading model...
Vocabulary size: 7643
Input tensor shape: torch.Size([1, 1, 32, 163])
------------------------------------------------------------
OCR RESULT:
------------------------------------------------------------
拍江买
------------------------------------------------------------
Inference time on NPU: 0.1862s
Done.
运行精度对比脚本:
python3 compare_cpu_npu.py --image /path/to/your/image.png该脚本会:
| 指标 | 值 |
|---|---|
| CPU 推理耗时 | 0.0750s |
| NPU 推理耗时 | 0.1810s |
| CPU 识别文本 | 拍江买 |
| NPU 识别文本 | 拍江买 |
| 文本完全匹配 | 是 |
| Logits 最大绝对误差 | 0.09243870 |
| Logits 平均绝对误差 | 0.01525087 |
| Logits MSE | 0.00035933 |
| 概率最大绝对误差 | 0.01134920 |
| 概率平均绝对误差 | 0.00000080 |
| Argmax 位置一致性 | 100.00% |
| 总体概率误差 | 0.0001% |
精度测试结论:NPU与CPU推理误差为0.0001%,符合精度误差小于1%的要求。
# 加载模型
from inference import ViTSTR, load_vocab, preprocess_image, ctc_decode
# 加载模型和词汇表
model = ViTSTR(num_classes=7644)
checkpoint = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(checkpoint, strict=True)
model.eval()
# 加载词汇表
chars = load_vocab("vocab.txt")
# 预处理图像
image_tensor = preprocess_image("test_image.png")
# 推理
with torch.no_grad():
logits = model(image_tensor)
logits_seq = logits[:, 1:, :] # 移除 CLS token
# CTC 解码
text = ctc_decode(logits_seq[0], chars, blank=0)
print(f"Recognized: {text}")如需在 NPU 上运行,只需将模型移至 NPU:
import torch_npu
model = model.npu()
image_tensor = image_tensor.npu()| 设备 | 推理耗时 | 加速比 (NPU/CPU) |
|---|---|---|
| CPU | 0.0750s | 1.00× |
| NPU | 0.1810s | 0.41× |
注:NPU 推理耗时包含模型加载和数据传输开销,对于批量推理场景,NPU 可发挥更大优势。
本仓库提供完整的推理脚本,支持 CPU 和 NPU 双平台推理:
# NPU 推理
python3 inference.py --device npu
# CPU 推理
python3 inference.py --device cpu推理完成后会输出推理结果和耗时,表明模型在 NPU 上推理成功。
#+NPU #+CV #+OCR #+手写体识别 #+ConvNeXt #+ViT #+昇腾 #+PyTorch