c++ - 如何从 C++ 中的原始指针数据构造一个 tensorflow::Tensor

Question

我想在操作期间更改输出张量的底层存储。

我有一个新数据的原始指针(float*)。我想在启动内核并返回之前将输出张量设置为这个新数据，这样我就可以劫持这个操作。

但是我对什么时候应该删除原始指针感到困惑，因为张量构造似乎是一个浅拷贝。我只能在所有张量使用完毕后删除原始指针。但是我怎样才能收到通知呢？

Answer 1

在 TensorFlow 运行时中没有用于执行此操作的公共(public) API，但可以使用 C API 方法从原始指针创建 Tensor 对象 TF_NewTensor() ，具有以下签名:

// Return a new tensor that holds the bytes data[0,len-1].
//
// The data will be deallocated by a subsequent call to TF_DeleteTensor via:
//      (*deallocator)(data, len, deallocator_arg)
// Clients must provide a custom deallocator function so they can pass in
// memory managed by something like numpy.
extern TF_Tensor* TF_NewTensor(TF_DataType, const int64_t* dims, int num_dims,
                               void* data, size_t len,
                               void (*deallocator)(void* data, size_t len,
                                                   void* arg),
                               void* deallocator_arg);

在内部，这创建了一个引用计数的 TensorBuffer拥有原始指针所有权的对象。 (不幸的是，只有 C API 有 friend access 来直接从 TensorBuffer 创建一个 tensorflow::Tensor。这是一个 open issue。) deallocator 函数在引用计数降为零时使用 data、len 和 dellocator_arg 的值调用。