Tensorflow新Op CUDA内核内存管理

Question

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

现在我正在使用 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] ; 删除[] var)

Answer 1

对于此问题没有直接的公共(public)指南。我通常只是让 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. 见下文

考虑一下，为了简单起见，只需添加两个矩阵 ( full example )。TensorFlow-Operations 通常包含以下结构:

操作描述具有REGISTER_OP，负责形状检查并设置输出形状 ( example )

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

仿函数用于实现本身，例如

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 调用之间的数据(例如一次性索引结构)，请使用此选项。[ example ]
• allocate_temp 只是临时内存，在 Compute 方法生命周期结束时不会保留。 [example ]

但我强烈建议阅读 source-code here 中的评论然后根据您的用例做出决定。