在多会话场景中，如何在一个tf :: Session（）中隐藏GPU？

Question

我有一个C ++程序，该程序使用在单独的tensorflow::Session中运行的两个不同的分类器。其中一种模型需要在GPU上运行，而第二种模型却很小，我希望它可以在CPU上运行。

我尝试使用以下方法创建会话：

auto options = tensorflow::SessionOptions();
options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0);
m_session.reset(tensorflow::NewSession(options));

但这似乎只是触发了默认的“猪GPU”行为。

我也尝试过使用options.config.mutable_gpu_options()->set_visible_device_list("-1")，但是TF代码中有一个特定的检查列表中的无效设备，这会引发错误。

在运行程序之前在环境中设置CUDA_VISIBLE_DEVICES=-1当然不是一种选择，因为对于需要它的会话，GPU应该是可见的。

有人知道如何禁止一个会话使用GPU？

Answer 1

一个简单的解决方法是使用CUDA_VISIBLE_DEVICES=-1临时设置putenv("CUDA_VISIBLE_DEVICES=-1");并在创建会话后将其重置。

std::string str_tmp = "CUDA_VISIBLE_DEVICES=";
str_tmp += getenv("CUDA_VISIBLE_DEVICES");
putenv("CUDA_VISIBLE_DEVICES=-1");

#Create session

#Reset environment variable
putenv(str_tmp);

但是，可以有一种更清洁的方法来完成此操作而不更改您的环境变量。