我在 Tensorflow 中使用 GPU CUDA 内核实现了一个相当复杂的新操作。 这个操作需要大量动态内存分配变量，这些变量不是张量，并且在操作完成后被释放，更具体地说，它涉及使用哈希表。
现在我正在使用 cudaMalloc() 和 cudaFree()，但我注意到 Tensorflow 有自己的类型，称为 Eigen::GPUDevice，它能够在 GPU 上分配和释放内存。
我的问题：
Eigen::GPUDevice 管理GPU 内存是最佳实践吗？Eigen::GPUDevice 而不是 CUDA API，我可以“自动”启用多 GPU 支持，因为可以将不同的 GPUDevices 传递给 Op；CPUDevice 类型也管理内存而不是使用 C++ 语法（即 auto var = new int[100]; delete[] var）【问题讨论】：
标签： tensorflow gpu
没有针对此问题的直接公开指南。我通常只是让 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));
cudaMalloc 或 new type[num] 调用。考虑一下，为了简单起见，只需添加两个矩阵 (full example)。 TensorFlow-Operations 通常包含以下结构：
Op description有REGISTER_OP，负责检查形状，设置输出形状（example）
OpKernel 负责分配内存、获取指向输入的指针和设置内容（见上文或this）
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,...)
编辑
Compute 方法生命周期结束时不会保留的tmp 内存。 [example]但我强烈建议阅读source-code here 中的评论，然后根据您的用例决定。
【讨论】：
allocate_temp()而不是allocate_persistent() 正如mrry的回答所建议的那样吗？
allocate_output。 GPU 内存是一种稀有资源（对大多数人来说），所以我通常会释放我的东西。
int a[3] = {1,2,3}，我应该关心这个吗？
最佳做法是使用OpKernelContext::allocate_persistent() 方法以tensorflow::Tensor 的形式分配内存，该方法比对OpKernel::Compute() 的单个调用要长。它为设备使用适当的Allocator*，因此如果内核在 GPU 设备上运行，它将为该特定设备分配 GPU 内存，如果它在 CPU 设备上运行，它将分配 CPU 内存。
【讨论】：