如何通过PyO3将Rust的tch::Tensor转换为Python可识别类型？

问题背景 #

尝试用PyO3和Rust开发Python扩展包，希望在Rust函数中返回tch::Tensor类型给Python，但直接返回触发类型转换错误，手动包装结构体实现IntoPy trait也无法解决类型不匹配问题。

错误原因 #

PyO3无法直接识别tch::Tensor类型，因为它不是PyO3定义的Python对象；同时tch::Tensor默认没有实现PyO3所需的IntoPy或OkWrap trait，导致无法自动转换为Python可识别的对象。

解决方案 #

要实现tch::Tensor到Python对象的转换，需要借助tch-rs的Python交互特性，将Rust端的Tensor转换为Python端的torch.Tensor对象：

1. 启用tch的Python特性：在Cargo.toml中为tch依赖添加python特性，确保tch-rs支持与Python的Tensor互转。
2. 正确转换Tensor：使用tch::Tensor提供的to_py_tensor方法，将Rust Tensor转换为Python可识别的Py<PyAny>对象。
3. 简化错误处理：直接实现From<TchError> for PyErr，让PyO3自动处理tch错误到Python异常的转换。

完整代码示例 #

Cargo.toml

[package]
name = "pyo3_utils"
version = "0.1.0"
edition = "2021"

[dependencies]
pyo3 = { version = "0.20", features = ["extension-module"] }
tch = { version = "0.14", features = ["python"] }

src/lib.rs

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
use tch::{Tensor, Kind};
use tch::TchError;

// 实现TchError到PyErr的转换，简化错误处理
impl From<TchError> for PyErr {
    fn from(error: TchError) -> Self {
        PyValueError::new_err(format!("PyTorch操作错误: {}", error))
    }
}

#[pyfunction]
fn create_tensor(py: Python<'_>) -> PyResult<Py<PyAny>> {
    // 创建Rust端Tensor
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0])
        .to_kind(Kind::Float);
    
    // 转换为Python端torch.Tensor
    tensor.to_py_tensor(py)
}

#[pymodule]
fn pyo3_utils(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(create_tensor, m)?)?;
    Ok(())
}

代码说明 #

• to_py_tensor方法：当tch启用python特性后，Tensor会新增该方法，负责将Rust端的Tensor转换为Python环境中的torch.Tensor对象，返回PyResult<Py<PyAny>>方便错误处理。
• 错误转换：通过From<TchError> for PyErr的实现，PyO3会自动将tch的错误转换为Python的ValueError，并附带具体错误信息，便于Python端调试。
• 函数参数：create_tensor需要接收Python<'_>参数，用于获取当前Python解释器上下文，完成跨语言对象转换。

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