使用pybind11实现任意NumPy数组转std::vector的技术问题

问题描述 #

使用pybind11开发Python插件时，遇到NumPy数组转std::vector<float>的类型转换错误。Python端传入{"varname": numpy_array}格式的字典，尝试直接py::cast转换时触发RuntimeError: Unable to cast Python instance of type <class 'numpy.ndarray'> to C++ type 'std::vector<float, std::allocator<float> >'错误。

当前解析函数代码：

std::vector<float> Plg::process(py::dict &inputs)
{
  std::vector<std::string> keys;
  //len(data) == len(keys)
  std::vector<std::vector<std::vector<float>>> data; //!!! 无法正常工作
  for (auto item : inputs)
  {
    std::string k = py::cast<std::string>(item.first);
    keys.push_back(k);
    std::vector<std::vector<float>> d = py::cast<std::vector<std::vector<float>>>(item.second);
    data.push_back(d);
  }
  //_process(data); - 实际处理函数，此处用模拟结果代替:
  std::vector<float> results;
  results.push_back(10.2);
  return results;
}

已知py::array_t可以处理扁平化数组，但输入必须是带字符串键、任意维度NumPy数组的字典，目标是将任意维度（如(10,10,2)、(200,)、(2,)）的NumPy数组转为扁平的std::vector<float>，并寻求无需memcpy的优化方案。

解决方案 #

1. 任意维度NumPy数组转扁平std::vector #

不能直接用py::cast把NumPy数组转成嵌套vector，正确做法是先将字典值转为py::array_t<float>，再利用其接口实现扁平化：

修改后的处理函数示例：

std::vector<float> Plg::process(py::dict &inputs)
{
    std::vector<std::string> keys;
    std::vector<std::vector<float>> flattened_data;

    for (auto &item : inputs)
    {
        // 提取字典键
        std::string k = py::cast<std::string>(item.first);
        keys.push_back(k);

        // 将NumPy数组转为py::array_t<float>，自动处理类型和布局转换
        py::array_t<float> arr = py::cast<py::array_t<float>>(item.second);
        
        // 直接通过size获取总元素数，构造对应大小的vector
        std::vector<float> flat_vec(arr.size());
        // 利用memcpy批量复制连续内存（仅适用于C风格连续布局的数组）
        std::memcpy(flat_vec.data(), arr.data(), arr.size() * sizeof(float));
        
        flattened_data.push_back(flat_vec);
    }

    // 调用实际处理逻辑，此处保留模拟结果
    std::vector<float> results;
    results.push_back(10.2);
    return results;
}

如果数组是非连续布局，可通过arr.request()获取buffer信息后，用循环逐个赋值：

py::buffer_info buf = arr.request();
std::vector<float> flat_vec(buf.size);
float* ptr = static_cast<float*>(buf.ptr);
for (size_t i = 0; i < buf.size; ++i) {
    flat_vec[i] = ptr[i];
}

2. 无需memcpy的优化实现 #

• 直接复用NumPy内存：如果你的_process函数可以直接接收指针和长度参数，完全不需要复制数据，直接使用NumPy数组的内存：

  // 假设_process支持指针+长度的参数形式
  void _process(float* data_ptr, size_t data_size) {
      // 处理逻辑
  }
  
  // 在process函数中：
  py::array_t<float> arr = py::cast<py::array_t<float>>(item.second);
  py::buffer_info buf = arr.request();
  // 强制转为C风格连续布局（非连续时会自动复制，可根据需求调整）
  py::array_t<float> arr_contiguous = arr.reshape({buf.size}).view<py::array::c_style>();
  _process(static_cast<float*>(arr_contiguous.data()), arr_contiguous.size());
  

• 用迭代器构造vector：如果必须使用std::vector，可以利用pybind11提供的迭代器直接构造，代码更简洁，底层效率和memcpy相近：

  std::vector<float> flat_vec(arr.begin<float>(), arr.end<float>());
  

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