简洁版
- 在tensorflow/core/ops 目录下找注册的operation名称,即OpName
- 在tensorflow/core/kernels目录下找与operation名称一致的源码文件
或直接在项目文档中搜索如下关键字:
REGISTER_KERNEL_BUILDER(Name("OpName")
(注意将OpName替换为步骤1中的operation名称,结尾是一个右括号)
啰里啰嗦版
通过TF的在线API查找operation的文档说明,注意TF版本。确定operation的准确名称。
clone TF在 GitHub 上的源码,checkout对应版本分支。
TF在线的API文档本身就是按照代码module分块组织在一起的,大多数情况下,我们要查的某个operation的实现源码就在对应的module下面。例如:
- tf.image下的api都是关于图像处理的方法
- tf.math下的api都是关于数学计算的方法
- tf.nn下的api都是关于神经网络的方法
- tf.quantization下的api都是关于量化的方法
- tf.io下的api都是关于数据输入、输出的方法
所有这些modules都在tensorflow/core/ops目录下有对应的module-name_ops.cc,这个文件中注册了该module中的operations,我们可以在该文件中确认该operation在代码中的准确名称(类名)。所以,每个operation必然在某个module的注册文件中注册,注册方式是:
REGISTER_OP("OpName")
.Input("arg1_name: type1")
.Input("arg2_name: type2") ...
.Output("out_variable: type")
.Attr("")
...
这个文件只是负责注册module中的operation,并非实现。
要查找某个operation的实现,需要到tensorflow/core/kernels目录下找该operation对应的准确名称(类名)文件即可。
kernels目录下的文件就是operation的实现,一般operation的实现对应两种或更多种设备的优化实现,例如intel CPU、nvidia GPU CUDA、ARM、Google TPU的。注意区分