如何在tensorflow中在gpu和cpu之间拉/推送数据

How can I pull/push data between gpu and cpu in tensorflow

本文关键字：之间 数据 cpu tensorflow 中在 gpu      更新时间：2023-10-16

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

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

我通常在调试过程中返回张量本身，以便一个简单的sess.run获取结果。否则，显示数据的唯一方法是C++内部的传统printf。假设张量是自定义操作的输出，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;
相关文章：