【Pytorch】API之from_blob函数

torch::from_blob 是 PyTorch C++ API (LibTorch) 中的一个函数，用于从已有的内存块（Blob）直接创建一个张量（Tensor），​无需复制数据。这在需要高效处理已有内存数据（例如从其他库或硬件获取的数据）时非常有用。

​函数签名​

torch::Tensor torch::from_blob(
    void* data,                          // 内存指针
    at::IntArrayRef sizes,               // 张量形状（如 {3, 224, 224}）
    const torch::TensorOptions& options = torch::TensorOptions()  // 设备、数据类型等
);

​关键参数详解​

1. ​data

   • 指向已有内存块的指针（例如 float*、uint8_t*）。
   • 内存需由用户自行管理，确保在张量使用期间有效。

2. ​sizes

   • 张量的维度形状，例如 {batch_size, channels, height, width}。

3. ​options

   • 指定张量的属性，常用选项：
     • dtype: 数据类型（如 torch::kFloat32、torch::kInt64）。
     • device: 设备（如 torch::kCPU、torch::kCUDA）。
     • requires_grad: 是否启用梯度计算（默认为 false）。

​使用示例​

​示例 1：从 CPU 内存创建张量​

#include <torch/torch.h>

int main() {
    // 创建一个浮点数组（假设内存由其他库提供）
    float data[] = {1.0, 2.0, 3.0, 4.0};

    // 从内存创建 2x2 的张量
    torch::Tensor tensor = torch::from_blob(
        data,                           // 数据指针
        {2, 2},                         // 形状为 2x2
        torch::TensorOptions().dtype(torch::kFloat32)  // 数据类型
    );

    std::cout << tensor << std::endl;
    // 输出:
    // 1  2
    // 3  4

    return 0;
}

​示例 2：指定设备（GPU）​​

// 假设已有 GPU 内存指针 float* gpu_data;
torch::Tensor gpu_tensor = torch::from_blob(
    gpu_data,
    {3, 256, 256},
    torch::TensorOptions().device(torch::kCUDA).dtype(torch::kFloat32)
);

​注意事项​

1. ​内存所有权​

   • from_blob 不会接管内存所有权，需确保内存生命周期覆盖张量使用范围。
   • 若内存由 new 或 malloc 分配，需手动释放。

2. ​内存对齐​

   • 确保内存布局与张量预期格式一致（例如连续存储）。

3. ​设备一致性​

   • 如果数据在 GPU 上，需确保 options.device() 设置为 torch::kCUDA。

4. ​步长（Strides）​​

   • 默认步长基于连续内存，若数据非连续，需显式指定 strides 参数：

     torch::Tensor tensor = torch::from_blob(
         data, 
         {2, 2}, 
         /* strides= */ {1, 2},  // 自定义步长
         options
     );

​Python 中的等价操作​

在 Python 中，类似功能可通过 torch.frombuffer 或 torch.from_numpy 实现：

import numpy as np
import torch

# 从 NumPy 数组创建（共享内存）
data_np = np.array([1, 2, 3], dtype=np.float32)
tensor = torch.from_numpy(data_np)

# 从缓冲区创建（需明确数据类型）
data_bytes = b"\x00\x00\x80\x3f"  # 小端序的 1.0
tensor = torch.frombuffer(data_bytes, dtype=torch.float32)

​常见问题​

• ​Q: 为什么修改原数据会改变张量？​​
  A: from_blob 是零拷贝操作，张量与原数据共享内存。

• ​Q: 如何避免悬垂指针？​​
  A: 确保原内存的生命周期长于张量，或使用 torch::Tensor::clone() 复制数据。

• ​Q: 是否支持多线程？​​
  A: 需要确保对内存的读写操作是线程安全的。