Tensorflow C++ 快速 Tensor 深拷贝

Question

I am writing a custom tensorflow op in C++ and I want to know how I might perform a deep copy of one tensor into another. In other words, I want an element-wise copy from one tensor to another such that they are not made to share an underlying memory buffer.

The closest thing I've found is DeepCopy (defined in tensor_util.h ). The problem is I need this operation to be fast and the documentation states clearly that this function has not been optimized for speed:

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

A bit more detail: I am using Tensorflow on the GPU, so basically what I want to do is to have TF initiate a cudaMemcpy where both source and target are device pointers. TF does give access to tensor pointers (tensor.tensor_data().data()), but you run into problems pretty quick if you try to cudaMemcpy with them (they have to do with overlapping cuda contexts; not nice).

Thanks in advance!

Answer 1

Currently for "fast" deep copies, TensorFlow uses Eigen to implement CPU and GPU versions. Being Eigen, the code is deeply templated, but the relevant code is here (copied from 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;
  }
};