# export.py —— 在训练环境里跑一次 import torch model = MyModel().eval() # 必须 eval()：关掉 dropout / BN 的训练行为 # 输入名字要与 Rust 侧喂数据的名字一致，建议起可读的名字 input_ids = torch.zeros((1, 128), dtype=torch.long) attention_mask = torch.ones((1, 128), dtype=torch.long) # 动态轴声明：不写这些，导出来后 batch/seq 会被写死成 1 和 128 # 这是"明明 Python 里能跑，Rust 里换了个长度就报错"的根因 dynamic_axes = { "input_ids": {0: "batch", 1: "seq_len"}, "attention_mask": {0: "batch", 1: "seq_len"}, "logits": {0: "batch", 1: "seq_len"}, } torch.onnx.export( model, (input_ids, attention_mask), "model.onnx", input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes=dynamic_axes, opset_version=17, # 选一个 ORT 支持的 opset，不要一味追新 do_constant_folding=True, ) # 导完立刻自检：用 onnx 库跑一遍形状推断，比等到 Rust 里报错强 import onnx m = onnx.load("model.onnx") onnx.checker.check_model(m) for i in m.graph.input: print("input:", i.name, [d.dim_param or d.dim_value for d in i.type.tensor_type.shape.dim]) for o in m.graph.output: print("output:", o.name, [d.dim_param or d.dim_value for d in o.type.tensor_type.shape.dim])

# 文本分类 / 特征提取模型的标准导出方式 pip install optimum[exporters] optimum-cli export onnx --model bert-base-chinese --task text-classification onnx/bert/ # 导出目录里会有：model.onnx、tokenizer.json、config.json # Rust 侧只需要前两个文件 # 想让模型更小更快：导出后再做量化（PyTorch 侧动态量化，不需要校准数据） from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic("model.onnx", "model.int8.onnx", weight_type=QuantType.QInt8) # 体积降到 1/4，CPU 上通常快 2 倍左右，精度损失一般 < 1%

[dependencies] # download-binaries：编译期自动下载对应平台的 onnxruntime 预编译包，最省事 # 生产环境若要用系统里已有的 ORT（共用一份、便于统一升级），改用 load-dynamic ort = { version = "2", features = ["download-binaries", "ndarray"] } # ort = { version = "2", features = ["load-dynamic", "ndarray"] } ndarray = "0.16" # 三个 feature 的作用： # download-binaries —— 构建时下载 onnxruntime 静态/动态库并链接 # load-dynamic —— 不链接，运行时通过 dlopen 加载，路径由 ORT_DYLIB_PATH 指定 # ndarray —— 提供 TensorRef::from_array_view 这类与 ndarray 互转的便利方法

use ort::session::builder::{GraphOptimizationLevel, SessionBuilder}; use ort::session::Session; pub fn build_session(model_path: &str, intra_threads: usize) -> ort::Result<Session> { // 进程级初始化：可以设置日志级别，排查 ORT 内部问题时会用到 // ort::init().with_name("order-api").commit()?; let mut builder = Session::builder()? // 1) 图优化级别：Level3 = 全部优化（算子融合、常量折叠、布局改写） // 用 Level1 会让推理明显变慢，除非在排查"优化后结果不对"的问题 .with_optimization_level(GraphOptimizationLevel::Level3)? // 2) 算子内部并行线程数：一次推理内部用几个核 // 注意它是【线程池】，不是"每请求一线程"；设太大反而抢占 .with_intra_threads(intra_threads)? // 3) 同时跑几个算子（层间并行）：通常保持默认 .with_inter_threads(1)? // 4) 内存模式：开启后复用缓冲区，减少分配；排查内存问题时可以关掉对比 .with_memory_pattern(true)? // 5) 执行提供者（GPU 等），见后面章节 // .with_execution_providers([...])? // 6) 允许把优化后的模型缓存到磁盘，跳过重复优化 .with_optimized_model_path(&format!("{model_path}.opt"))?; let session = builder.commit_from_file(model_path)?; Ok(session) } // 加载完先把输入输出的名字与形状打出来 —— 这是排查一切"喂不进去"的第一步 pub fn dump_io(session: &Session) { for input in session.inputs() { println!("input name={} dtype={:?}", input.name(), input.dtype()); } for output in session.outputs() { println!("output name={} dtype={:?}", output.name(), output.dtype()); } // 还有一类信息容易被忽略：模型自带的元数据（训练时的超参、类别标签） if let Ok(meta) = session.metadata() { // 很多团队会把 id2label 塞在 metadata 里，部署时直接读，不用另传文件 println!("producer={:?}", meta.producer()); } }

use ndarray::{Array2, Array3, Array4}; use ort::value::TensorRef; // 文本模型：两个 long 类型的输入，形状 [batch, seq_len] let input_ids: Array2<i64> = Array2::from_shape_vec((1, 5), vec![101, 872, 1962, 102, 0])?; let attention_mask: Array2<i64> = Array2::from_shape_vec((1, 5), vec![1, 1, 1, 1, 0])?; // TensorRef::from_array_view 是【借用】—— 不拷贝数据、不拥有内存 // 注意生命周期：底层 ndarray 必须在 run 期间一直存活 let outputs = session.run(ort::inputs![ // 键名必须与模型里的 input 名字完全一致（区分大小写） "input_ids" => TensorRef::from_array_view(&input_ids)?, "attention_mask" => TensorRef::from_array_view(&attention_mask)?, ])?; // 图像模型：NCHW 还是 NHWC，取决于导出时模型的约定 // ResNet 等 torchvision 模型是 NCHW：[batch, channels, height, width] let image_nchw: Array4<f32> = Array4::zeros((1, 3, 224, 224)); // 部分 TF / TFLite 转换来的模型是 NHWC：[batch, height, width, channels] let image_nhwc: Array4<f32> = Array4::zeros((1, 224, 224, 3)); // 需要有拥有权的 Tensor（数据要活过 ndarray 的作用域）时用 from_array（会拷贝） // let owned = Tensor::from_array(([1usize, 3, 224, 224], vec![0f32; 3 * 224 * 224]))?;

use ndarray::ArrayViewD; // 1) 按名字取，比按索引取可读得多（模型输出顺序变了也不会错） let logits: ArrayViewD<f32> = outputs["logits"].try_extract_array::<f32>()?; // 形状一般是 [batch, num_class]（分类）或 [batch, seq_len, vocab]（语言模型） // 先断言形状，把"我以为的"变成"编译期/运行期会检查的" let shape = logits.shape(); assert_eq!(shape.len(), 2, "期望二维 logits，实际 {:?}", shape); // 2) Softmax：把 logit 变成概率。记得减去最大值防止 exp 溢出 fn softmax(logits: &[f32]) -> Vec<f32> { let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max); let exps: Vec<f32> = logits.iter().map(|x| (x - max).exp()).collect(); let sum: f32 = exps.iter().sum(); exps.into_iter().map(|e| e / sum).collect() } // 3) argmax：取最大概率的类别 fn argmax(scores: &[f32]) -> usize { scores.iter().enumerate() .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) .map(|(i, _)| i).unwrap_or(0) } // 4) top-k：推荐/召回场景需要前 k 个而不是只要最大那个 fn top_k(scores: &[f32], k: usize) -> Vec<(usize, f32)> { let mut idx: Vec<usize> = (0..scores.len()).collect(); // 部分排序：K 远小于 N 时比全排序快得多 idx.select_nth_unstable_by(k.min(scores.len().saturating_sub(1)), |a, b| { scores[*b].partial_cmp(&scores[*a]).unwrap_or(std::cmp::Ordering::Equal) }); idx[..k.min(scores.len())].iter().map(|&i| (i, scores[i])).collect() } // 用起来：取 batch 里第一条样本 let row = logits.index_axis(ndarray::Axis(0), 0); // 注意：标准布局下 as_slice() 一定成功；不标准时再 to_vec() 兜底 let row_vec: Vec<f32> = row.iter().copied().collect(); let probs = softmax(&row_vec); let best = argmax(&probs); println!("预测类别 = {}, 置信度 = {:.4}", best, probs[best]);

use tokio::sync::{mpsc, oneshot}; struct Job { input_ids: Vec<i64>, // 每条请求自带一个"回信地址"，批处理完成后逐个应答 reply: oneshot::Sender<anyhow::Result<usize>>, } /// 批处理 worker：单线程持有 session（因为 run 需要 &mut self） async fn batch_worker( mut rx: mpsc::Receiver<Job>, mut session: Session, max_batch: usize, max_wait: std::time::Duration, ) { loop { // 1) 阻塞等第一条，等不到就退出（通道关闭） let Some(first) = rx.recv().await else { break }; let mut batch = vec![first]; // 2) 在 max_wait 内尽量多攒几条，攒到 max_batch 就立刻发 // 这两个参数直接决定"吞吐 vs 延迟"的平衡点 let deadline = tokio::time::Instant::now() + max_wait; while batch.len() < max_batch { match tokio::time::timeout_at(deadline, rx.recv()).await { Ok(Some(job)) => batch.push(job), // 超时或通道关闭，立即发车 _ => break, } } // 3) padding 到同一长度：动态 shape 模型可以不做，静态模型必须做 let seq_len = batch.iter().map(|j| j.input_ids.len()).max().unwrap_or(1); let rows = batch.len(); let mut flat = vec![0i64; rows * seq_len]; // 0 作为 pad token let mut mask = vec![0i64; rows * seq_len]; for (r, job) in batch.iter().enumerate() { for (c, &id) in job.input_ids.iter().enumerate() { flat[r * seq_len + c] = id; mask[r * seq_len + c] = 1; // 真实 token 标 1，pad 保持 0 } } // 4) 一次 run 处理整批 let ids: ndarray::Array2<i64> = ndarray::Array2::from_shape_vec((rows, seq_len), flat).unwrap(); let m: ndarray::Array2<i64> = ndarray::Array2::from_shape_vec((rows, seq_len), mask).unwrap(); let out = session.run(ort::inputs![ "input_ids" => ort::value::TensorRef::from_array_view(&ids).unwrap(), "attention_mask" => ort::value::TensorRef::from_array_view(&m).unwrap(), ]); // 5) 把结果按行拆开，逐条回信（每条的 argmax 独立算） match out { Ok(outputs) => { let logits = outputs["logits"].try_extract_array::<f32>().unwrap().to_owned(); for (r, job) in batch.into_iter().enumerate() { let row = logits.index_axis(ndarray::Axis(0), r); let row_vec: Vec<f32> = row.iter().copied().collect(); let _ = job.reply.send(Ok(argmax(&row_vec))); } } Err(e) => { // 单批失败要逐条报错，不能让调用方永久挂住 let msg = e.to_string(); for job in batch { let _ = job.reply.send(Err(anyhow::anyhow!(msg.clone()))); } } } } }

use ort::session::builder::SessionBuilder; // macOS：CoreML（走 Apple Neural Engine / GPU） #[cfg(target_os = "macos")] fn with_accel(b: SessionBuilder) -> ort::Result<SessionBuilder> { use ort::ep::CoreML; // 顺序很重要：先列的优先使用，失败会自动回退到后面的 // 所以常规写法是 [GPU_EP, CPU EP]，ORT 自带 CPU 兜底 b.with_execution_providers([ CoreML::default() // 把整个图交给 ANE 通常会失败回退，按算子子图划分反而更稳 .with_model_format(ort::ep::coreml::ModelFormat::MLProgram) .build(), ]) } // Linux + NVIDIA：CUDA #[cfg(target_os = "linux")] fn with_accel(b: SessionBuilder) -> ort::Result<SessionBuilder> { use ort::ep::CUDA; b.with_execution_providers([ CUDA::default() .with_device_id(0) // fp16 推理：显存减半、速度提升，精度损失通常可接受 // 但必须在精度敏感的模型上验证，别默认打开 // .with_fp16(true) .build(), ]) } // Windows 上可以用 DirectML 走任意显卡；Intel CPU 服务器用 OpenVINO // 注意 CUDA/CoreML 的 EP 需要 onnxruntime 构建时带了对应支持， // 官方预编译包的 GPU 版本要单独下载（CPU 版不含 CUDA） // 怎么确认 EP 真的生效了？打开 ORT 日志看它选了哪个 EP，别只看代码 // ort::init().with_log_level(ort::logging::LogLevel::Info).commit()?;

# 1) 动态库到底在哪？链接到了谁？ ldd target/release/order-api | grep onnx # Linux otool -L target/release/order-api | grep onnx # macOS # 2) 动态加载模式：确认路径生效（程序启动时打印一行日志最省事） ORT_DYLIB_PATH=/usr/local/lib/libonnxruntime.dylib ./order-api # 在代码里验证：println!("{:?}", ort::info()); # 3) 版本匹配：ORT 动态库版本与 ort crate 期望的 ABI 版本不一致也会加载失败 # 动态库与 crate 的版本对应关系写在 ort 的 README / MINOR_VERSION 常量里 python3 -c "import onnxruntime; print(onnxruntime.__version__)" # 容器里别忘了把这个库一起 COPY 进最终镜像 # FROM scratch 的极简镜像尤其容易漏 —— 它连 glibc 都没有