美文网首页
如何定位operation在TF源码中的位置

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

作者: Aspirinrin | 来源:发表于2019-12-27 09:36 被阅读0次

    简洁版

    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下面。例如:

    • 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的。注意区分

    相关文章

      网友评论

          本文标题:如何定位operation在TF源码中的位置

          本文链接:https://www.haomeiwen.com/subject/igmroctx.html