TensorFlow 自定义分配器和从 Tensor 访问数据

Posted 2023-02-17

【中文标题】TensorFlow 自定义分配器和从 Tensor 访问数据

【英文标题】：TensorFlow Custom Allocator and Accessing Data from Tensor

【发布时间】：2016-09-30 17:38:03

【问题描述】：

在 TensorFlow 中，您可以出于各种原因创建自定义分配器（我正在为新硬件做这件事）。由于设备的结构，我需要使用一些元素的结构作为我的数据指针，分配器将其返回为void*。

在我正在编写的内核中，我可以访问张量，但我不需要获取我编写的指针结构。检查类，似乎我可以通过 tensor_t.buf_->data() 得到这个结构

Tensor::buf_

TensorBuffer::data()

问题是我找不到执行此操作的代码，我担心它不安全（很有可能！）或者有更标准的方法来执行此操作。

有人可以确认这是一个好主意还是坏主意？并提供替代方案（如果存在）？

【参考方案1】：

您也可以使用Tensor::tensor_data().data() 来访问原始指针，而无需通过 DMAHelper 使用奇怪的间接。

【讨论】：

如何从原始指针构造张量？我什么时候应该删除原始指针？谢了。

【参考方案2】：

四天后...

void* GetBase(const Tensor* src) 
  return const_cast<void*>(DMAHelper::base(src));

来自GPUUtils

DMAHelper::base() 是一个友元类方法，它可以使用私有 Tensor::base() 来获取数据指针。

implementation 表明这只是我在另一个abstraction 之后想要做的事情的包装。我猜这是一种更安全的获取指针的方法，应该改用它。