原创文章,转载请注明: 转载自慢慢的回味
本文链接地址: 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 国际许可协议进行许可。