在多GPU模式下的tensorflow c ++ SetDefaultDevice

时间:2018-09-06 15:23:15

标签: c++ tensorflow machine-learning deep-learning

我想在多个GPU上加载同一图形进行推理,但是我无法使用graph :: SetDefaultDevice将图形与设备关联。该问题不在SetDefaultDevice中发生,而是在以后与图形创建会话时发生。这是从tensorflow的example_trainer.cc中提取的一个简单示例

#include <tensorflow/core/platform/env.h>
#include <tensorflow/core/public/session.h>
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/graph/default_device.h"

int main() {
  using namespace tensorflow;
  using namespace tensorflow::ops;
  Scope root = Scope::NewRootScope();
  auto A = Const(root, { {3.f, 2.f}, {-1.f, 0.f} }); 
  auto b = Const(root, { {3.f, 5.f} }); 
  auto v = MatMul(root.WithOpName("v"), A, b, MatMul::TransposeB(true));

  GraphDef def;
  TF_CHECK_OK(root.ToGraphDef(&def));

  graph::SetDefaultDevice(false ? "/device:GPU:0" : "/cpu:0", &def);
  /*
  for (auto &node: *def.mutable_node()) {
        node.set_device("/cpu:0");
        std::cout << node.name() << " = '" << node.device() <<"'"<< std::endl;
  }
  std::cout << "=======================\n";
  */
  SessionOptions options;
  std::unique_ptr<Session> session(NewSession(options));
  TF_CHECK_OK(session->Create(def));
  return 0;
}

运行时出现以下错误

2018-09-06 18:18:13.853316: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2018-09-06 18:18:13.856079: F /home/daniel/tensorflow_cc/example/example.cpp:27] Non-OK-status: session->Create(def) status: Not found: No attr named '/cpu:0' in NodeDef:
     [[Node: Const = Const[dtype=DT_FLOAT, value=Tensor<type: float shape: [2,2] values: [3 2][-1]...>, _device="/cpu:0"]()]]
     [[Node: Const = Const[dtype=DT_FLOAT, value=Tensor<type: float shape: [2,2] values: [3 2][-1]...>, _device="/cpu:0"]()]]
Aborted (core dumped)

如果我删除SetDefault Device调用,它将运行完美。我也尝试在具有GPU的计算机上执行此操作,但是它没有用。

我知道问题出在SetDefaultDevice上,因为手动创建每个节点的设备最终在创建会话时会遇到相同的问题。

Const = '/cpu:0'
Const_1 = '/cpu:0'
v = '/cpu:0'
=======================
2018-09-06 18:15:05.966337: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2018-09-06 18:15:05.969048: F /home/daniel/tensorflow_cc/example/example.cpp:26] Non-OK-status: session->Create(def) status: Not found: No attr named '/cpu:0' in NodeDef:
     [[Node: Const = Const[dtype=DT_FLOAT, value=Tensor<type: float shape: [2,2] values: [3 2][-1]...>, _device="/cpu:0"]()]]
     [[Node: Const = Const[dtype=DT_FLOAT, value=Tensor<type: float shape: [2,2] values: [3 2][-1]...>, _device="/cpu:0"]()]]
Aborted (core dumped)

1 个答案:

答案 0 :(得分:1)

这似乎仅对于单片构建(--config = monolithic)才是问题,即在构建libtensorflow_cc.so时。我不确定,但可能与

有关

https://github.com/tensorflow/tensorflow/issues/5379 https://github.com/tensorflow/tensorflow/issues/16291