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

Question

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

Answer 1

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

我通常在调试过程中返回张量本身，以便简单的sess.run获取结果。否则，显示数据的唯一方法是传统的printf内部 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;