如何在pybind11中加速array_t构建？C++转PyTorch张量优化

你的核心问题在于现有方法都做了不必要的内存拷贝，导致耗时过高。以下是两种零拷贝的优化方案，能大幅降低数据传输耗时：

方案一：零拷贝构造py::array_t，再转PyTorch Tensor #

直接利用std::vector的内存构造py::array_t，避免拷贝；Python端用torch.from_numpy直接复用numpy数组的内存（同样零拷贝）。

C++代码修改： #

int main(){
    const int length = 500000;
    float list[length];
    // 填充list数据...
    std::vector<float> v(list, list+length);

    // 零拷贝构造py::array_t，关联vector的内存
    py::array_t<float> args(
        {length},                  // 数组形状
        {sizeof(float)},           // 元素步长
        v.data(),                  // 指向vector数据的指针
        py::cast(v)                // 所有权绑定，确保vector在array_t存活时不被释放
    );

    py::module_ nd_to_tensor = py::module_::import("inference");
    py::object result = nd_to_tensor.attr("inference")(args);
}

Python代码修改： #

import torch

def inference(arr):
    # torch.from_numpy直接复用arr的内存，无拷贝
    tensor = torch.from_numpy(arr)
    # 后续推理逻辑...
    return tensor

方案二：C++直接构造PyTorch Tensor，直接传给Python #

如果你的项目可以链接libtorch，直接在C++端创建PyTorch Tensor，跳过中间的py::array_t，实现端到端零拷贝。

C++代码修改： #

#include <torch/torch.h>
#include <pybind11/pybind11.h>

namespace py = pybind11;

int main(){
    const int length = 500000;
    float list[length];
    // 填充list数据...
    std::vector<float> v(list, list+length);

    // 从vector内存构造Tensor，零拷贝
    torch::Tensor tensor = torch::from_blob(
        v.data(),
        {length},                  // Tensor形状
        torch::kFloat32            // 数据类型
    );

    // 将Tensor转为py::object传给Python
    py::object py_tensor = py::cast(tensor);

    py::module_ nd_to_tensor = py::module_::import("inference");
    py::object result = nd_to_tensor.attr("inference")(py_tensor);
}

Python代码修改： #

import torch

def inference(tensor):
    # 直接使用传入的Tensor进行推理，无任何拷贝
    output = model(tensor)
    return output

原方法耗时高的原因 #

• 方法一：默认的py::cast(v)会创建一个新的py::array_t并拷贝std::vector的所有数据，50万float的拷贝耗时约0.5秒。
• 方法二：绑定的Vector在Python中是一个自定义序列类型，torch.Tensor(Vector)会遍历序列逐个拷贝元素，因此耗时更长（1.4秒）。

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