tensorflow c++ op 生成 python调用接口

迈不过友情╰ 2022-11-11 13:14 217阅读 0赞

## 例如 ##

from tensorflow.python.ops import gen_math_ops

看一下bazel的调用顺序

*  生成op的lib  
    生成 math\_ops\_op\_lib, 根据文件的名字，不要带.cc 后缀， 生成library, 生成的library 叫做 math\_ops\_op\_lib

tf_gen_op_libs(
       is_external = False,
       op_lib_names = [
       	 "math_ops",
    ])

*  生成python的接口

tf_gen_op_wrapper_private_py(
        name = "math_ops_gen",  # 结尾必须是_gen
        visibility = [
            "//learning/brain/google/python/ops:__pkg__",
            "//learning/brain/python/ops:__pkg__",
            "//tensorflow/compiler/tests:__pkg__",
            "//tensorflow/contrib/quantization:__pkg__",
            "//tensorflow/python/kernel_tests:__pkg__",
        ],
    )

到这里就可以使用from tensorflow.python.ops import gen\_math\_ops

## 具体分析 ##

tf\_gen\_op\_wrapper\_private\_py

def tf_gen_op_wrapper_private_py(
            name,
            out = None,
            deps = [],
            require_shape_functions = True,
            visibility = []):
        if not name.endswith("_gen"):
            fail("name must end in _gen") # 结尾必须是_gen
        if not visibility:
            visibility = ["//visibility:private"]
        bare_op_name = name[:-4]  # Strip off the _gen
        # 提取 name 为 math_ops
        tf_gen_op_wrapper_py(
            name = bare_op_name,
            out = out,
            visibility = visibility,
            deps = deps,
            require_shape_functions = require_shape_functions,
            generated_target_name = name,
            api_def_srcs = [
                "//tensorflow/core/api_def:base_api_def",
                "//tensorflow/core/api_def:python_api_def",
            ],
        )

tf\_gen\_op\_wrapper\_py

def tf_gen_op_wrapper_py(
            name,
            out = None,
            hidden = None,
            visibility = None,
            deps = [],
            require_shape_functions = False,
            hidden_file = None,
            generated_target_name = None,
            op_whitelist = [],
            cc_linkopts = [],
            api_def_srcs = []):
        if (hidden or hidden_file) and op_whitelist:
            fail("Cannot pass specify both hidden and op_whitelist.")
    
        # Construct a cc_binary containing the specified ops.
        tool_name = "gen_" + name + "_py_wrappers_cc"
        if not deps:
            deps = [str(Label("//tensorflow/core:" + name + "_op_lib"))]   # 没有制定op的c++ library, 则会加上 math_ops_op_lib 
        # 生成gen_math_op_py_wrappers_cc的二进制
        tf_cc_binary(
            name = tool_name,
            copts = tf_copts(),
            linkopts = if_not_windows(["-lm", "-Wl,-ldl"]) + cc_linkopts,
            linkstatic = 1,  # Faster to link this one-time-use binary dynamically
            visibility = [clean_dep("//tensorflow:internal")],
            deps = ([
                clean_dep("//tensorflow/core:framework"),
                clean_dep("//tensorflow/python:python_op_gen_main"),   # 调用python_op_gen_main 生成 python文件
            ] + deps),
        )
    
        # 通过c++ 的binary 生成 python文件
        if not out:
            out = "ops/gen_" + name + ".py"
    
        if hidden:
            op_list_arg = ",".join(hidden)
            op_list_is_whitelist = False
        elif op_whitelist:
            op_list_arg = ",".join(op_whitelist)
            op_list_is_whitelist = True
        else:
            op_list_arg = "''"
            op_list_is_whitelist = False
    
        # Prepare ApiDef directories to pass to the genrule.
        if not api_def_srcs:
            api_def_args_str = ","
        else:
            api_def_args = []
            for api_def_src in api_def_srcs:
                # Add directory of the first ApiDef source to args.
                # We are assuming all ApiDefs in a single api_def_src are in the
                # same directory.
                api_def_args.append(
                    "$$(dirname $$(echo $(locations " + api_def_src + ") | cut -d\" \" -f1))",
                )
            api_def_args_str = ",".join(api_def_args)
    
        if hidden_file:
            # `hidden_file` is file containing a list of op names to be hidden in the
            # generated module.
            native.genrule(
                name = name + "_pygenrule",
                outs = [out],
                srcs = api_def_srcs + [hidden_file],
                tools = [tool_name] + tf_binary_additional_srcs(),
                cmd = ("$(location " + tool_name + ") " + api_def_args_str +
                       " @$(location " + hidden_file + ") " +
                       ("1" if require_shape_functions else "0") + " > $@"),
            )
        else:
            native.genrule(
                name = name + "_pygenrule",
                outs = [out],
                srcs = api_def_srcs,
                tools = [tool_name] + tf_binary_additional_srcs(),
                cmd = ("$(location " + tool_name + ") " + api_def_args_str + " " +
                       op_list_arg + " " +
                       ("1" if require_shape_functions else "0") + " " +
                       ("1" if op_list_is_whitelist else "0") + " > $@"),
            )
    
        # Make a py_library out of the generated python file.
        if not generated_target_name:
            generated_target_name = name
        native.py_library(
            name = generated_target_name,
            srcs = [out],
            srcs_version = "PY2AND3",
            visibility = visibility,
            deps = [
                clean_dep("//tensorflow/python:framework_for_generated_wrappers_v2"),
            ],
            # Instruct build_cleaner to try to avoid using this rule; typically ops
            # creators will provide their own tf_custom_op_py_library based target
            # that wraps this one.
            tags = ["avoid_dep"],
        )