feat(model): 新增 ModelLoader 链式构建 API 及 ORT GPU/Tract 多线程配置

- 在 ddddocr-core 中定义 ModelBuilder Trait 及其错误类型
- ddddocr-ort 支持 use_gpu、device_id 及 num_threads 链式配置与 CUDA 硬件加速
- ddddocr-tract 基于 multithread-mm 特性支持 CPU 线程数控制
- 支持基于 tract-linalg 配置推理线程数,显式引入 tract-linalg 的 multithread-mm 特性,控制 GEMM 算子并发
- 优化线程池加载策略,适配 Tokio 异步及 CLI 等多场景
This commit is contained in:
2026-07-27 20:22:24 +08:00
parent 44dae08221
commit 7d159c5702
28 changed files with 1583 additions and 70 deletions

View File

@@ -0,0 +1,114 @@
use crate::loader::Error;
use crate::loader::error::{BuildError, ParseError, Result};
use crate::types::Session;
use ddddocr_core::ModelBuilder;
use ort::session::Session as OrtSession;
use ort::session::builder::SessionBuilder;
use std::sync::{Arc, Mutex};
pub struct OrtModelLoader;
impl OrtModelLoader {
/// 获取针对 ORT 后端的链式构建器
pub fn builder() -> OrtModelBuilder {
OrtModelBuilder::default()
}
}
/// ORT 专用的链式构建器
#[derive(Debug, Clone)]
pub struct OrtModelBuilder {
use_gpu: bool,
device_id: i32,
intra_threads: Option<usize>,
}
impl Default for OrtModelBuilder {
fn default() -> Self {
Self {
use_gpu: false,
device_id: 0,
intra_threads: None,
}
}
}
impl OrtModelBuilder {
/// 开启或关闭 GPU 加速
pub fn use_gpu(mut self, enable: bool) -> Self {
self.use_gpu = enable;
self
}
/// 指定 GPU 设备 ID
pub fn device_id(mut self, id: i32) -> Self {
self.device_id = id;
self
}
pub fn num_threads(mut self, threads: usize) -> Self {
self.intra_threads = Some(threads);
self
}
/// 内部辅助方法:根据当前的配置构建 ORT 底层的 SessionBuilder
fn create_session_builder(&self) -> Result<SessionBuilder> {
let mut builder = OrtSession::builder().map_err(|e| BuildError::BuildFailed(e))?;
// 如果用户显式设置了线程数,则配置给 ORT
if let Some(threads) = self.intra_threads {
builder = builder
.with_intra_threads(threads)
.map_err(|e| BuildError::Threads(format!("设置线程数失败: {e}")))?;
}
if self.use_gpu {
// 根据 ort 库版本配置 CUDA 执行提供者 (Execution Provider)
#[cfg(feature = "cuda")]
{
use ort::ep::CUDAExecutionProvider;
let cuda_ep = CUDAExecutionProvider::default().with_device_id(self.device_id);
builder = builder
.with_execution_providers([cuda_ep.build()])
.map_err(|e| {
BuildError::EnabledCudaFailed(format!("配置 CUDA 硬件加速失败: {e}"))
})?
}
#[cfg(not(feature = "cuda"))]
{
// 如果用户明确开启了 GPU但 Feature 没编译进去,明确抛错提醒
return Err(BuildError::NotEnabledCuda(
"未启用 CUDA 支持:请在 Cargo.toml 中为 ddddocr-ort 开启 `cuda` feature"
.to_string(),
))?;
}
}
Ok(builder)
}
}
impl ModelBuilder for OrtModelBuilder {
type Session = Session;
type Error = Error;
fn model_for_path<P>(&self, model_path: P) -> Result<Session>
where
P: AsRef<std::path::Path>,
{
let path_ref = model_path.as_ref();
let mut builder = self.create_session_builder()?;
// Session::builder() 会返回 Result<SessionBuilder, ort::Error>
let session = builder
.commit_from_file(path_ref)
.map_err(|e| ParseError::Path(path_ref.display().to_string(), e))?;
Ok(Arc::new(Mutex::new(session))) // 这里的session需要包装下
}
/// 策略 B从内存字节流加载模型配合 include_bytes! 使用)
fn model_from_bytes(&self, model_bytes: &[u8]) -> Result<Session> {
let mut builder = self.create_session_builder()?;
let session = builder
.commit_from_memory(model_bytes)
.map_err(ParseError::Bytes)?;
Ok(Arc::new(Mutex::new(session)))
}
}