跳转至

使用其他推理引擎

引言

因为版面分析模型输入图像尺寸固定，故可使用onnxruntime-gpu来提速。因为rapid_layout库默认依赖是CPU版onnxruntime，如果想要使用GPU推理，需要手动安装onnxruntime-gpu。详细使用和评测可参见AI Studio

使用GPU

pip install rapid_layout

# 请确保 onnxruntime-gpu 与当前 GPU/CUDA 版本对应
# 参见 https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements
pip install onnxruntime-gpu

使用

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
from rapid_layout import EngineType, ModelType, RapidLayout, RapidLayoutInput

cfg = RapidLayoutInput(
    model_type=ModelType.PP_LAYOUT_CDLA,
    engine_type=EngineType.ONNXRUNTIME,
    engine_cfg={"use_cuda": True, "cuda_ep_cfg.gpu_id": 1},
)
layout_engine = RapidLayout(cfg=cfg)

img_path = "https://raw.githubusercontent.com/RapidAI/RapidLayout/718b60e927ab893c2fad67c98f753b2105a6f421/tests/test_files/layout.jpg"
results = layout_engine(img_path)
print(results)

results.vis("layout_res.png")

使用NPU

详细配置参数参见：engine_cfg.yaml

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
from rapid_layout import EngineType, ModelType, RapidLayout, RapidLayoutInput

cfg = RapidLayoutInput(
    model_type=ModelType.PP_LAYOUT_CDLA,
    engine_type=EngineType.ONNXRUNTIME,
    engine_cfg={"use_cann": True, "cann_ep_cfg.gpu_id": 0},
)
layout_engine = RapidLayout(cfg=cfg)

img_path = "https://raw.githubusercontent.com/RapidAI/RapidLayout/718b60e927ab893c2fad67c98f753b2105a6f421/tests/test_files/layout.jpg"
results = layout_engine(img_path)
print(results)

results.vis("layout_res.png")

评论