TensorFlow C 快速张量深副本

Tensorflow C++ fast Tensor deep copy

本文关键字：张量深 副本 TensorFlow      更新时间：2023-10-16

我正在编写C 中的自定义TensorFlow OP，我想知道如何将一个张量的深副本执行到另一个张量。换句话说，我想要一个元素从一个张量到另一种张量的元素，以免它们共享基础的内存缓冲区。

我发现的最接近的是deepcopy（在tensor_util.h中定义）。问题是我需要此操作很快，并且文档清楚地指出，该功能尚未针对速度进行优化：

// DeepCopy returns a tensor whose contents are a deep copy of the
// contents of 'other'.  This function is intended only for
// convenience, not speed.

更多的细节：我在GPU上使用Tensorflow，因此我想做的是让TF启动一个cudamemcpy，其中源和目标都是设备指针。TF确实可以访问张量指针（tensor.tensor_data（）。data（）），但是如果您尝试与它们进行cudamemcpy（它们必须与重叠的cuda上下文有关；不好）。

预先感谢！

当前用于"快速"深副本，TensorFlow使用eigen实现CPU和GPU版本。作为特征，代码是深层的模板，但是相关代码在这里（从dense_update_ops.h复制）：

template <typename Device, typename T>
struct DenseUpdate<Device, T, ASSIGN> {
  void operator()(const Device& d, typename TTypes<T>::Flat params,
                  typename TTypes<T>::ConstFlat update) {
    params.device(d) = update;
  }
};
相关文章：