c++ - 在多 session 场景中，如何从一个 tf::Session() 中隐藏 GPU？

Question

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

我尝试使用以下方法创建 session :

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 应该对需要它的 session 可见。

有谁知道如何禁止只有一个一个 session 使用 GPU？

Answer 1

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

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);

但是，可能有一种更简洁的方法可以在不更改环境变量的情况下执行此操作。