Question

我有一个训练模型并使用python tensorflow保存。现在我想在c ++应用程序中执行预测。

我正在加载这样的模型（在Tensorflow 1.3.1上）：

std::unique_ptr<tensorflow::SavedModelBundle> bundle(new tensorflow::SavedModelBundle());

tensorflow::Status status = tensorflow::LoadSavedModel(
  tensorflow::SessionOptions(),
  tensorflow::RunOptions(),
  modelPath,
  { tensorflow::kSavedModelTagServe },
  bundle.get());

如何指定运行此模型的设备？具体来说，我希望能够强制模型在cpu而不是gpu上运行。

目前，我这样做是通过设置环境变量CUDA_VISIBLE_DEVICES = -1，但这不允许我例如在gpu上有一个模型，在cpu上有另一个模型。

谢谢！

Answer 1

会话选项允许您配置会话中可见的设备，因此您可以阻止创建GPU设备，如下所示：

tensorflow::SessionOptions options;
auto* device_count = options.config.mutable_device_count();
device_count->insert({"CPU", 1});
device_count->insert({"GPU", 0});

然后在options的调用中使用tensorflow::SessionOptions()代替LoadSavedModel()。

Answer 2

您可以在会话选项中设置以下内容，但这也是会话范围：

tensorflow::SessionOptions options = TF::SessionOptions();
options.config.mutable_gpu_options()->force_gpu_compatible();

如果你想特别为每个模型或变量提供它，你应该在构建模型时使用with说明符，然后在将它序列化为protobuf图时存储：

with tf.device('/cpu:0'):
    a = tf.constant([1.0, 2.0], shape=[2, 1])

或：

with tf.device('/device:GPU:0'):
    b = tf.constant([1.0, 2.0], shape=[2, 1])

如何指定用于保存模型的设备

2 个答案: