TensorFlow图的构建

原创文章，转载请注明： 转载自慢慢的回味

本文链接地址: TensorFlow图的构建

本文就来分析tensorflow图的构建，以及运行图的设备和线程池的创建。TensorFlow使用数据流图将计算表示为独立的指令之间的依赖关系，数据流是一种用于并行计算的常用编程模型。在数据流图中，节点表示计算单元，边缘表示计算使用或产生的数据。

数据流可以为TensorFlow提供多项优势：

  • 并行处理：通过使用明确的边缘来表示操作之间的依赖关系，系统可以轻松识别能够并行执行的操作。
  • 分布式执行：通过使用明确的边缘来表示操作之间流动的值，TensorFlow 可以将您的程序划分到连接至不同机器的多台设备上（CPU、GPU 和 TPU）。
  • 编译：TensorFlow 的 XLA 编译器可以使用数据流图中的信息生成更快的代码。
  • 可移植性：数据流图是一种不依赖于语言的模型代码表示法。您可以使用 Python 构建数据流图，将其存储在 SavedModel 中，并使用 C++ 程序进行恢复，从而实现低延迟的推理。


回目录
类TensorfowTest.cc中有如下代码，它调用c api完成一个图的创建：

	TF_Graph* graph = TF_NewGraph();

在文件c_api.cc中有TF_Graph::TF_Graph构造函数：
graph的类型为Graph，refiner的类型为ShapeRefiner。

TF_Graph::TF_Graph()
    : graph(tensorflow::OpRegistry::Global()),
      refiner(graph.versions().producer(), graph.op_registry()),
      delete_requested(false),
      parent(nullptr),
      parent_inputs(nullptr) {
  // Tell the shape refiner to also run shape inference on functions.
  refiner.set_function_library_for_shape_inference(&graph.flib_def());
}
 
TF_Graph* TF_NewGraph() { return new TF_Graph; }


在graph.cc文件中，构造函数Graph::Graph创建了一个底层的Graph实例。由于是空图，所以需要添加一个开始节点_SOURCE即后面的root节点和一个最终节点_SINK。

Graph::Graph(const OpRegistryInterface* ops)
    : ops_(ops, FunctionDefLibrary()),
      versions_(new VersionDef),
      arena_(8 << 10 /* 8kB */) {
  versions_->set_producer(TF_GRAPH_DEF_VERSION);
  versions_->set_min_consumer(TF_GRAPH_DEF_VERSION_MIN_CONSUMER);
 
  // Initialize the name interning table for assigned_device_name.
  device_names_.push_back("");
  DCHECK_EQ(0, InternDeviceName(""));
 
  // Source and sink have no endpoints, just control edges.
  NodeDef def;
  def.set_name("_SOURCE");
  def.set_op("NoOp");
  Status status;
  Node* source = AddNode(def, &status);
  TF_CHECK_OK(status);
  CHECK_EQ(source->id(), kSourceId);
 
  def.set_name("_SINK");
  Node* sink = AddNode(def, &status);
  TF_CHECK_OK(status);
  CHECK_EQ(sink->id(), kSinkId);
 
  AddControlEdge(source, sink);
}

在graph.cc文件中，Graph::AddNode方法完成节点的创建。先通过ops_.LookUpOpDef查找op类型的定义，然后通过构造参数NodeProperties完成Node的实例化。

Node* Graph::AddNode(NodeDef node_def, Status* status) {
  const OpDef* op_def;
  status->Update(ops_.LookUpOpDef(node_def.op(), &op_def));
  if (!status->ok()) return nullptr;
 
  DataTypeVector inputs;
  DataTypeVector outputs;
  status->Update(InOutTypesForNode(node_def, *op_def, &inputs, &outputs));
  if (!status->ok()) {
    *status = AttachDef(*status, node_def);
    return nullptr;
  }
 
  Node* node = AllocateNode(std::make_shared<NodeProperties>(
                                op_def, std::move(node_def), inputs, outputs),
                            nullptr);
  return node;
}

通过shape_refiner.cc中的ShapeRefiner::ShapeRefiner和graph_runner.cc中的GraphRunner::GraphRunner生成graph_runner_对象，graph_runner_负责创建线程运行设备：

ShapeRefiner::ShapeRefiner(int graph_def_version,
                           const OpRegistryInterface* ops)
    : graph_def_version_(graph_def_version),
      ops_registry_(ops),
      graph_runner_(Env::Default()) {}
 
GraphRunner::GraphRunner(Env* env)
    : device_deleter_(NewSingleThreadedCpuDevice(env)),
      device_(device_deleter_.get()) {}

文件single_threaded_cpu_device.cc中创建一个单线程的线程池，这个线程只会处理一些简单的图的计算，Eigen是一个矩阵运算的库：

Device* NewSingleThreadedCpuDevice(Env* env) {
  return new SingleThreadedCpuDevice(env);
}
class SingleThreadedCpuDevice : public Device {
 public:
  explicit SingleThreadedCpuDevice(Env* env)
      : Device(env, Device::BuildDeviceAttributes("/device:CPU:0", DEVICE_CPU,
                                                  Bytes(256 << 20),
                                                  DeviceLocality())) {
    eigen_worker_threads_.num_threads = kNumThreads;
    eigen_worker_threads_.workers = GraphRunnerThreadPool();
    eigen_device_.reset(new Eigen::ThreadPoolDevice(
        eigen_worker_threads_.workers->AsEigenThreadPool(),
        eigen_worker_threads_.num_threads));
    set_tensorflow_cpu_worker_threads(&eigen_worker_threads_);
    set_eigen_cpu_device(eigen_device_.get());
  }

本作品采用知识共享署名 4.0 国际许可协议进行许可。

发表回复