使用 c++，我如何知道张量流张量是在 cuda 还是 cpu 中？答案

【问题标题】：using c++, How can I know whether a tensorflow tensor is in cuda or cpu?使用 c++，我如何知道张量流张量是在 cuda 还是 cpu 中？
【发布时间】：2022-07-29 13:51:14
【问题描述】：

我正在写一个基于tensorflow的模型推理引擎，有没有知道cuda设备或cpu上的TF Tensor？

【问题讨论】：

标签： tensorflow

【解决方案1】：

刚刚发现可以使用cudaPointerGetAttributes来检查张量是否在CUDA上，详情附后。顺便说一句，以下代码来自：https://github.com/triton-inference-server/tensorflow_backend/blob/main/src/tensorflow_backend_tf.cc#L403

void
TensorImpl::Init()
{
  nonstring_base_ = nullptr;
  nonstring_byte_size_ = 0;
  gpu_tensor_ = false;

  // Implement differently for string and non-string
  if (tftensor_.dtype() != tensorflow::DT_STRING) {
    auto flat = tftensor_.bit_casted_shaped<char, 1>(
        {tftensor_.NumElements() *
         tensorflow::DataTypeSize(tftensor_.dtype())});
    nonstring_base_ = static_cast<char*>(flat.data());
    nonstring_byte_size_ = flat.size();

    cudaPointerAttributes attributes;
    cudaError_t err = cudaPointerGetAttributes(&attributes, nonstring_base_);
    gpu_tensor_ =
        ((err == cudaSuccess) && (attributes.type == cudaMemoryTypeDevice));
  }
}

【讨论】：