如何将原始指针投射到特定形状的 pytorch 张量？

Question

我从 C++ 库中获得一个原始指针，我想将其解释为（以“

reinterpret_cast

”式的方式）为特定形状的 pytorch 张量。由于代码是在性能关键部分执行的，所以我真的想要以确保不执行堆分配和/或复制操作。

这是我现在得到的：

def as_tensor(pointer, shape):
    return torch.from_numpy(numpy.array(numpy.ctypeslib.as_array(pointer, shape = shape)))

shape = (2, 3, 4)
x = torch.zeros(shape)

p = ctypes.cast(x.data_ptr(), ctypes.POINTER(ctypes.c_float))
y = as_tensor(p, shape)

之前真的有必要转换为 numpy 数组吗？而且我也不能 100% 确定对

numpy.array(...)

的调用是否不会复制

as_array()

调用所指向的内容。

Answer 1

要将 C++ 库中的原始指针解释为特定形状的 PyTorch 张量而不执行堆分配或复制操作，您可以通过指定数据指针和形状直接创建 PyTorch 张量。这是代码的修改版本，旨在避免不必要的副本：

import torch
import ctypes

def as_tensor(pointer, shape):
    # Create a PyTorch tensor from a raw pointer without copying data
    tensor = torch.as_tensor(pointer, dtype=torch.float32)
    # Reshape the tensor to the specified shape
    return tensor.view(shape)

shape = (2, 3, 4)
x = torch.zeros(shape)

p = ctypes.cast(x.data_ptr(), ctypes.POINTER(ctypes.c_float))
y = as_tensor(p, shape)

在此修订版中：

我们使用

torch.as_tensor()

直接从原始指针创建 PyTorch 张量，而不涉及任何复制操作。 通过使用

torch.as_tensor()

而不是 torch.from_numpy() ，我们避免了转换为 NumPy 数组的中间步骤，这可能涉及不必要的复制。

view()

方法用于将张量重塑为指定形状，而不更改其基础数据。 通过使用

torch.as_tensor()

，您可以以高性能的方式将原始指针解释为 PyTorch 张量，而无需承担额外内存分配或数据副本的风险。