当前位置: 代码迷 >> 综合 >> TensorFlow源码分析(4):OpKernelContruction类的构造函数
  详细解决方案

TensorFlow源码分析(4):OpKernelContruction类的构造函数

热度:91   发布时间:2023-09-07 08:14:08.0

在文章(2)中,介绍了ConstantOp类的构造函数。构造函数中唯一的一个参数是OpKernelConstruction类型的对象。这个对象包含了操作类需要的所有输入,参数以及运行环境。它定义在tensorflow/core/framework/op_kernel.h文件中。OpKernelConstruction类的成员函数比较多,无法在一篇文章中一一列举。本文先从它的构造函数入手,分析它的成员变量。

 

OpKernelConstruction类的构造函数定义在tensorflow/core/framework/op_kernel.cc文件中。

OpKernelConstruction::OpKernelConstruction(DeviceType device_type, DeviceBase* device, Allocator* allocator,const NodeDef* node_def, const OpDef* op_def, FunctionLibraryRuntime* flib,const DataTypeSlice& input_types, const MemoryTypeSlice& input_memory_types,const DataTypeSlice& output_types,const MemoryTypeSlice& output_memory_types, int graph_def_version,Status* status): device_type_(std::move(device_type)),device_(device),allocator_(allocator),def_(node_def),op_def_(op_def),flib_(flib),input_types_(input_types),input_memory_types_(input_memory_types),output_types_(output_types),output_memory_types_(output_memory_types),graph_def_version_(graph_def_version),status_(status) {}
  • 从上面的代码,可以清晰看到OpKernelConstruction类的成员变量。
  • device_type_和device_是指定操作运行的设备。这个设备可能是CPU和GPU两种,它们都可能是本地设备,或者远程设备。
  • allocator_是一个存储资源是分配器。如果操作运行在CPU上,它代表内存分配器;如果运行在GPU上,代表显存分配器。
  • def_是一个NodeDef类型的对象。我们已经说过,TensorFlow中所有的操作,在计算图中都表示成一个阶段。def_对象对应的就是当前ConstantOp操作在计算图对应的节点。
  • op_def_是OpDef类型的对象。在TensorFlow中,任何一个操作都有两部分组成。一个是OpDef类型的对象,它类似于C语言的声明。通过OpDef对象,我们可以告诉TensorFlow当前操作的输入、输出以及参数的类型。另外一个,就是我们重点研究的操作类。它相当于C语言的函数定义,用来实现操作的业务逻辑。那么前面的NodeDef对象又是做什么呢?它是用来控制输入与输出的传递。在Nodef对象中,会指定当前操作的输入是由哪些操作的输出给定;当前操作是输出,又将作为哪些操作的输入。
  • input_types用来指定各个输入的Tensor的元素的数据类型。input_types是一个数组,其大小与输入Tensor的个数相同。input_memory_types用来指定各个输入的Tensor分别存放在哪个存储设备上。当前,TensorFlow只支持内存和显存两种设备。
  • output_types和output_memory_types指定的是输出的数据类型和存储设备。
  • status_用于指定计算是否出错,并保存错误信息。

另外两个变量flib_和graph_def_version_作者没有深入研究,日后再补充。

  相关解决方案