Tensorflow新的Op CUDA内核内存管理

Question

我已经使用GPU CUDA内核在Tensorflow中实现了一个相当复杂的新Op.这个Op需要大量动态内存分配变量,这些变量不是张量的,并且在op完成后被释放,更具体地说它涉及使用散列表.

现在我正在使用cudaMalloc(),cudaFree()但我注意到Tensorflow有自己的类型Eigen::GPUDevice,它具有在GPU上分配和释放内存的能力.

我的问题:

1. Eigen::GPUDevice用于管理GPU内存是最佳做法;
2. 通过使用Eigen::GPUDevice而不是CUDA API,我"自动"启用多GPU支持,因为不同的GPUDevices可以传递给Op;
3. 我应该将这个想法扩展到CPU内核,看看是否有一种CPUDevice类型也管理内存而不是使用C++语法(即auto var = new int[100]; delete[] var)

Answer 1

这个问题没有直接的公共指南.我通常只是让TensorFlow分配这些信息

template<typename Device, typename Dtype>
class MyOp: public OpKernel {
{
public:
  explicit MyOp(OpKernelConstruction *context) :
      OpKernel(context)
  {
    // ...
  }

  void Compute(OpKernelContext *context) override
  {
    Tensor* tmp_var = nullptr;
    Tensor* output = nullptr;

    TensorShape some_shape, some_shape2;

    // temparily use this space
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_FLOAT, some_shape, &tmp_var));
    // allocate memory for output tensor
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, some_shape2, &output));

1. 无论需要什么内存,都应该由TensorFlow上下文分配,而不是通过自定义cudaMalloc或new type[num]调用.
2. 上下文应该为分配器提供信息
3. 见下文

考虑到,为简单起见,只需添加两个矩阵(完整示例).TensorFlow-Operations通常包含以下结构:

运算描述具有REGISTER_OP,这是负责形状检查,并设置输出形状(例如)

OpKernel负责分配内存,获取指向输入和设置内容的指针,(见上文或此内容)

Functor用于实现本身,就像

Tensor* output = nullptr;
Tensor* tmp_var = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output));
OP_REQUIRES_OK(ctx, ctx->allocate_temp(0, some_shape, &tmp_var));
// the function does not need to care about the memory allocation as everything is already setup at this point
::tensorflow::functor::MyFunctor<Device, Dtype>()(ctx, inputA, inputB, tmp_var, output);

你刚刚离开了

    // gpu version
    template <typename Dtype>
    struct MyFunctor<GPUDevice, Dtype> {
      void operator ()(::tensorflow::OpKernelContext* ctx,...)

    // cpu version
    template <typename Dtype>
    struct MyFunctor<CPUDevice, Dtype> {
      void operator ()(::tensorflow::OpKernelContext* ctx,...)

编辑

• allocate_persistent:如果你需要像一次性索引结构这样的Op调用之间的数据,请使用它.[ 例子 ]
• allocate_temp只是tmp内存,它不会在Compute方法生命周期结束时保留.[ 例子 ]

但我强烈建议在这里阅读源代码中的注释,然后根据您的用例决定.