如何使用内存地址访问Pytorch张量数据？

Question

我正在创建一个 pytorch 张量

import torch
x = torch.randn(100,100,device='cuda')
adr = x.data_ptr() #gives address to first element of x, e.g. 30119297024

有没有办法在只知道的情况下访问张量数据

adr

并做类似的事情

y = get_tensor_from_address(adr)

其中

y

和

x

指的是同一个张量？

Answer 1

您可以使用

ctypes

库访问该元素；然而，正如另一条评论中提到的，这种方法并不安全。

这里有两种访问元素的替代方法

1. 通过地址访问张量元素

   import torch
   import ctypes
   
   tensor = torch.rand((2, 3), dtype=torch.float32)
   data_ptr = tensor.data_ptr()
   
   row, col = 1, 0
   element_index = row * tensor.size(1) + col
   element_size = tensor.element_size()
   element_address = data_ptr + element_index * element_size
   
   element_value = ctypes.c_float.from_address(element_address).value
   assert element_value == tensor[row][col]
   

2. 将

   data_ptr

   投射到

   ctypes._Pointer

   import torch
   import ctypes
   
   tensor = torch.rand(2, 3, dtype=torch.float32)
   data_ptr = tensor.data_ptr()
   c_float_ptr = ctypes.cast(data_ptr, ctypes.POINTER(ctypes.c_float))
   
   row, col = 1, 0
   index = row * tensor.size(1) + col
   value = c_float_ptr[index]
   assert c_float_ptr[index] == tensor[row][col]
   

顺便说一句，如果你想访问 GPU 上的张量，只需将其复制到 CPU 即可。