TensorFlow C++ 初始化 Tensor 内存 到GPU 内存

最近使用TensorFlow C++版本实现神经网络的部署，我通过GPU 处理得到网络的输入值，因此输入值在GPU内存上保存， TF 输入tensor 的调用语句为 Tensor inputTensor(DT_FLOAT, TensorShape({ 1,2,3,1 }));  默认构造是将内存放到CPU上的。为了实现GPU 到GPU 的内存拷贝，而不是 GPU 到 Cpu 在从CPU 到GPU（通过PCIE总线内存拷贝耗时高），我们需要将inputTensor 内存初始化到GPU上通过实现下面代码就可以实现。

#include "tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h"
#include "tensorflow/core/common_runtime/gpu/gpu_cudamalloc_allocator.h"

tensorflow::GPUBFCAllocator * allocator = new tensorflow::GPUBFCAllocator(0,sizeof(float)* Col_num * tempfftsize);
    //tensorflow::Allocator* allocator = new AllocatorWrapper(0, tempfftsize * Col_num * sizeof(float));
    tensorflow::GPUcudaMallocAllocator *gpu_allocator = new tensorflow::GPUcudaMallocAllocator(gpu_allocator, 0);
    tensorflow::Tensor inputTensor(gpu_allocator,DT_FLOAT, tensorflow::TensorShape({ 1,Col_num,tempfftsize,1 }));
    auto inputTensor_flat = inputTensor.flat<float>();
    cudaMemcpy(&inputTensor_flat(0), d_LogSpec, tempfftsize * Col_num * sizeof(float), cudaMemcpyDeviceToDevice);//d_LogSpec为输入的GPU内存地址

更详细的介绍参考 https://github.com/tensorflow/tensorflow/issues/19283