Question

这个问题是我之前question的后续问题，我问过如何在TensorFlow操作系统中分配GPU或CPU内存。

我希望能够在TensorFlow Op中为GPU / CPU中的任何C ++类型分配内存。使用context可以分配新的张量（在CPU或GPU中），如果它们是＆＃34;常规＆＃34; C ++类型一可以获得C ++指针。例如：

    Tensor tensor;
    OP_REQUIRES_OK(context, context->allocate_temp(DT_FLOAT, TensorShape({n_elements}), &tensor));
    float * ptr = tensor.flat<float>().data();

我的问题是在尝试分配用户定义的类型时：

struct A{
    float a;
    int b;
};

然后没有相应的Tensorflow类型作为分配。我知道在这种情况下，我可以只分配float张量和int张量，但是对于更复杂的数据结构，这会使代码变得非常混乱。

我尝试通过分配字节（DT_UINT8）并使用reinterpret_cast来转换为所需的指针类型（作为我之前提出的问题中的答案中的一条评论）。这是我的实施：

template<typename t> void allocate(void ** ptr_address, int num_elements){
    Tensor tensor;
    int number_of_bytes = num_elements * sizeof(t);
    OP_REQUIRES_OK(context, context->allocate_temp(DT_UINT8, TensorShape({number_of_bytes}), &tensor));
    *ptr_address = reinterpret_cast<t*>(tensor.flat<unsigned char>().data());
}

理论上，我可以将A* a; allocate<A>(&a, 5);用于上面定义的类，而不是自动a = new A[5];或cudaMalloc((void**)&a, 5*sizeof(A))。

问题是在Op运行之后它会冻结而TensorFlow不会返回到python（我确保它通过printf调试运行）。我确信这是因为这种分配方法，因为当我使用常规C ++ / CUDA分配时，我不会遇到这个问题。

这种分配方法有问题还是其他问题？

在TensorFlow Op中分配C ++ / CUDA内存

0 个答案: