c - TensorFlow C 后端选择 GPU

Question

我已经使用 C 后端编写了一个完整的 TF 推理管道。目前，我正在开发具有多个 GPU (x8) 的硬件。它在 CPU 上运行良好，而不是在 GPU 上运行良好，因为我无法正确选择设备。

工作流程如下:单个线程已从保存的模型设置 session

TF_LoadSessionFromSavedModel(...)

然后，池中的线程为 C 后端执行通常的工作流程(设置输入/输出并运行)

 TF_NewTensor(...) // allocate input 
 TF_AllocateTensor(.....) // allocate ouput 
 TF_SessionRun(....) 

目前，我知道我想在哪个设备上执行我的代码，所以我正在使用 CUDA 驱动程序 API cudaSetDevice但是，它没有任何影响(默认情况下它始终在设备 0 上，请使用 nvidia-smi 检查)。如果我使用 CUDA_VISIBLE_DEVICES 强制设备我可以有效地选择其他设备 ID，但是 CUDA_VISIBLE_DEVICE=0,1,2,3,4,5,6,7结合 cudaSetDevice不起作用。

我怀疑 TF 在内部强制设备，也许可以使用 TF_SetConfig 来实现灵活性, 或 run_options的 TF_SessionRun .但是，C 后端的文档不存在。因此，如果 TF 向导在这里，我将不胜感激正确设置设备以执行 TF_SessionRun 的建议。 .

Answer 1

经过与 Google TF 开发人员的长期交流以及漫长的一天编码之后，我将回答我的问题。出发点:目前无法使用保存的模型制作多 GPU。为了成功，我给出了待办事项 list :

将保存的模型转换为旧方式卡住缓冲区(protobuf)。因此，图形和元数据将被整合到一个"file"中。进行这种转换的最佳方法是 python。本帖提供所有info .我们获得了一个 Protobuf。

阅读 C 中的 Protobufs 并获得 TF_graph*。我为您提供了 C 代码来做到这一点:

inline void deallocate_buffer(void *data, size_t) { 
    std::free(data);
}

inline TF_Buffer *read_buffer_from_file(const char *file) {
    const auto f = std::fopen(file, "rb");
    if (f == nullptr) {
        return nullptr;
    }

   std::fseek(f, 0, SEEK_END);
   const auto fsize = ftell(f);
   std::fseek(f, 0, SEEK_SET);

   if (fsize < 1) {
      std::fclose(f);
      return nullptr;
   }

   const auto data = std::malloc(fsize);
   std::fread(data, fsize, 1, f);
   std::fclose(f);

   TF_Buffer *buf = TF_NewBuffer();
   buf->data = data;
   buf->length = fsize;
   buf->data_deallocator = deallocate_buffer;

   return buf;
}

inline TF_Graph *load_graph_def(const char *file) {
    if (file == nullptr) {
        return nullptr;
    }

   TF_Buffer *buffer = read_buffer_from_file(file);
   if (buffer == nullptr) {
       return nullptr;
   }

   TF_Graph *graph = TF_NewGraph();
   TF_Status *status = TF_NewStatus();
   TF_ImportGraphDefOptions *opts = TF_NewImportGraphDefOptions();

   TF_GraphImportGraphDef(graph, buffer, opts, status);
   TF_DeleteImportGraphDefOptions(opts);
   TF_DeleteBuffer(buffer);

   check_status(status);
   TF_DeleteStatus(status);
   return graph;
}

将图形转换为 graphdef 并修改每个节点的执行设备。我提供主要元素。

TF_GraphToGraphDef(...)
// prepare the option for the new graph
std::string device = "/device:GPU:" + std::to_string(i); // i is an int
TF_ImportGraphDefOptions* graph_options = ...
TF_ImportGraphDefOptionsSetDefaultDevice(graph_options, device.c_str());
// create the new graph with correct device
TF_Graph *ngraph = TF_NewGraph();
TF_GraphImportGraphDef(ngraph, buffer, graph_options, status);
check_status(status);
// create the new session
TF_Session *session = TF_NewSession(ngraph, session_opts, status);

重要提示:必须为每个 GPU 设备创建一个新 session 。由您使用您最喜欢的线程系统来管理它。

最后，TF_NewSession 的 session_opts 必须使用十六进制 protobuf 的一些黑魔法正确设置。 (GPU ID 和软放置)

// this cryptic buffer is generated by
// import tensorflow as tf
// gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=0.5,
//                             allow_growth=True,visible_device_list='0,1,2,3,4,5,6,7')
//               config = tf.compat.v1.ConfigProto(gpu_options=gpu_options,allow_soft_placement=True )
//               serialized = config.SerializeToString()
//               print(list(map(hex, serialized)))

std::vector<uint8_t> config = {0x32, 0x1c, 0x9,  0x0,  0x0,  0x0,  0x0,  0x0,  0x0,  0xe0, 0x3f,
                               0x20, 0x1,  0x2a, 0xf,  0x30, 0x2c, 0x31, 0x2c, 0x32, 0x2c, 0x33,
                               0x2c, 0x34, 0x2c, 0x35, 0x2c, 0x36, 0x2c, 0x37, 0x38, 0x1};

TF_SetConfig(session_opts, config.data(), config.size(), status);

我希望这会帮助那些必须在 C 中运行他们的 TF 的人。