能否在Rust中导入Python训练的机器学习.pkl模型？

.pkl是Python专属的序列化格式，Rust没有原生支持，给你几个可行的解决思路：

方案1：将模型转换为跨语言格式（推荐） #

把.pkl模型导出为ONNX格式，这是机器学习领域的跨语言标准格式，Rust有成熟的库支持加载。

步骤： #

• 用Python将pkl模型转ONNX：

  import joblib
  from skl2onnx import convert_sklearn
  from skl2onnx.common.data_types import FloatTensorType
  
  # 加载pkl模型
  model = joblib.load("your_model.pkl")
  # 定义输入张量类型（根据你的模型输入调整）
  initial_type = [("float_input", FloatTensorType([None, 4]))]
  # 转换为ONNX
  onnx_model = convert_sklearn(model, initial_types=initial_type)
  # 保存ONNX模型
  with open("your_model.onnx", "wb") as f:
      f.write(onnx_model.SerializeToString())
  

  注：如果是TensorFlow/PyTorch模型，直接用框架自带的ONNX导出API即可

• Rust中加载ONNX模型：
  添加依赖到Cargo.toml：

  [dependencies]
  onnxruntime = { version = "0.18", features = ["cpu"] }
  ndarray = "0.15"
  

  加载代码示例：

  use onnxruntime::{environment::Environment, GraphOptimizationLevel, Session};
  use ndarray::Array;
  
  fn main() -> Result<(), Box<dyn std::error::Error>> {
      let env = Environment::builder()
          .with_name("test")
          .build()?;
  
      let session = Session::builder()?
          .with_optimization_level(GraphOptimizationLevel::Basic)
          .with_number_threads(1)
          .load_model(env, "your_model.onnx")?;
  
      // 构造输入数据（根据模型输入维度调整）
      let input_data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).into_shape((1, 4))?;
      let outputs = session.run(vec![input_data.into()])?;
  
      // 处理输出结果
      let output = outputs[0].try_extract::<Array<f32, _>>()?;
      println!("模型输出: {:?}", output);
  
      Ok(())
  }
  

方案2：通过Python绑定在Rust中调用Python加载模型 #

如果不想转换模型，可以用pyo3库在Rust中直接调用Python代码加载.pkl模型。

步骤： #

• 添加依赖到Cargo.toml：

  [dependencies]
  pyo3 = { version = "0.20", features = ["auto-initialize"] }
  

• Rust代码示例：

  use pyo3::prelude::*;
  
  fn main() -> PyResult<()> {
      // 初始化Python环境
      Python::with_gil(|py| {
          // 导入joblib加载模型
          let joblib = py.import("joblib")?;
          let model = joblib.call_method1("load", ("your_model.pkl",))?;
  
          // 构造输入数据并预测
          let input = py.eval("[1.0, 2.0, 3.0, 4.0]", None, None)?;
          let prediction = model.call_method1("predict", (input,))?;
  
          println!("预测结果: {:?}", prediction.extract::<Vec<f32>>()?);
          Ok(())
      })
  }
  

  注：需要确保运行环境中安装了Python和对应的依赖（如joblib、scikit-learn等）

方案3：使用Rust原生ML库加载特定模型 #

如果你的模型是scikit-learn的，可以尝试rust-learn这类库，但它只支持部分模型类型，局限性较大，不推荐作为通用方案。

内容的提问来源于stack exchange，提问作者Arkadeep Nag