如何定位operation在TF源码中的位置

2019-12-27  本文已影响0人  Aspirinrin

简洁版

  1. tensorflow/core/ops 目录下找注册的operation名称,即OpName
  2. 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下面。例如:

所有这些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的。注意区分

上一篇 下一篇

猜你喜欢

热点阅读