如何在张量流中的 gpu 和 cpu 之间拉/推数据

Question

我使用一个临时张量将数据存储在我定制的基于 gpu 的操作中。出于调试目的，我想通过 C++ 中的传统 printf 打印此张量的数据。我怎样才能把这个基于 gpu 的张量拉到 cpu 上，然后打印它的内容。 非常感谢。

Answer 1

如果 temporary 你的意思是 allocate_temp 而不是 allocate_output，则无法在 python 端获取数据。

我通常在调试期间返回张量本身，以便简单的sess.run 获取结果。否则，显示数据的唯一方法是传统的printf inside C++。鉴于您的张量是您的自定义操作的输出，tf.Print 简化了进一步的调试。

例子：

Tensor temp_tensor;
OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_FLOAT, some.shape(), &temp_tensor));

float* host_memory = new float[some.NumElements()];
cudaMemcpy(host_memory, temp_tensor.flat<Dtype>().data(), some.NumElements() * sizeof(float), cudaMemcpyDeviceToHost);
std::cout << host_memory[0] << std::endl;
std::cout << host_memory[1] << std::endl;
std::cout << host_memory[2] << std::endl;
delete[] host_memory;