机器学习中的python调用c

时间:2022-03-02作者:klpeng分类:Python浏览:755评论:0

简介

前言,当查看tf python 代码时, 碰到一个函数实际由C++ 实现时,

pybind11

动手学深度学习框架(2) - python端如何调用c++的代码

pybind11 是一个轻量级的只包含头文件(header-only)的 c++ 库,用于将 c++ 代码暴露给 python 调用(反之亦可,但主要还是前者)。

#include <pybind11/pybind11.h>          // pybind11的头文件
// PYBIND11_MODULE 是一个宏,实现一个 Python 扩展模块
PYBIND11_MODULE(pydemo, m){             // 定义Python模块pydemo,之后在 Python 脚本里必须用这个名字才能 import。m其实是 pybind11::module 的一个实例对象,它只是个普通的变量,起什么名字都可以,但为了写起来方便,一般都用“m”。
    m.doc() = "pybind11 demo doc";      // 模块的说明文档
    m.def("add", [](int a, int b) -> int { return a + b; });    // def函数,传递一个 Python 函数名和 C++ 的函数、函数对象或者是 lambda 表达式
}                                       // Python模块定义结束

假设这个 C++ 源文件名是“pybind.cpp”,用 g++ 把它编译成在 Python 里调用的模块,生成一个大概这样的文件:pydemo.cpython-35m-x86_64-linux-gnu.so

g++ pybind.cpp               \                  #编译的源文件
   -std=c++11 -shared -fPIC   \                 #编译成动态库
  `python3 -m pybind11 --includes` \            #获得 pybind11 所在的包含路径,让 g++ 能够找得到头文件
  -o pydemo`python3-config --extension-suffix`  #生成的动态库名字,前面必须是源码里的模块名,而后面那部分则是 Python 要求的后缀名

之后就可以在python 中使用

import pydemo
x = pydemo.add(1,2)
print(x)

进阶:C++ 里的类也能够等价地转换到 Python 里面调用

自定义算子

一个Op可以接收一个或者多个输入Tensor,然后产生零个或者多个输出Tensor,分别利用Input和Output定义。在注册一个Op之后,就需要继承OpKernel,实现他的计算过程Compute函数,在Compute函数中,我们可以通过访问OpKernelContext来获得输入和输出信息。当我们需要申请新的内存地址时,可以通过OpKernelContext去申请TempTensor或者PersistentTensor。一般Tensorflow的Op都采用Eigen来操作Tensor

Adding a Ncw Op对于 TensorFlow,可以自定义 Operation,即如果现有的库没有涵盖你想要的操作, 你可以自己定制一个。为了使定制的 Op 能够兼容原有的库,你必须做以下工作:

  1. 在一个 C++ 文件中注册新 Op. Op 的注册与实现是相互独立的. 在其注册时描述了 Op 该如何执行. 例如, 注册 Op 时定义了 Op 的名字, 并指定了它的输入和输出.REGISTER_OP("ZeroOut") .Input("to_zero: int32") .Output("zeroed: int32") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(0)); return Status::OK(); });
  2. 使用 C++ 实现 Op. 每一个实现称之为一个 “kernel”, 可以存在多个 kernel, 以适配不同的架构 (CPU, GPU 等)或不同的输入/输出类型.
  3. bazel(tf编译工具) 会检索所有op 并创建一个 Python wrapper. 这个wrapper是创建 Op 的公开 API. 当注册 Op 时, 会自动生成一个默认 默认的包装器. 既可以直接使用默认包装器, 也可以添加一个新的包装器.
  4. (可选) 写一个函数计算 Op 的梯度,在Python 中注册.@ops.RegisterGradient("ZeroOut") def _zero_out_grad(op, grad): xxxxxxxxx
  5. (可选) 写一个函数, 描述 Op 的输入和输出 shape. 该函数能够允许从 Op 推断 shape.
  6. 测试 Op, 通常使用 Pyhton。如果你定义了梯度,你可以使用Python的GradientChecker来测试它。

There are two main mechanisms for op and kernel registration:

tf.load_op_library()

Bazel BUILD文件如下,执行 bazel build ${BAZEL_ARGS[@]} 可以得到 tensorflow/core/user_ops/zero:zero_out.so PS: 类似于执行了 上文中的g++

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
tf_custom_op_library(
    name = "zero_out.so",       #  target name
    srcs = ["zero_out.cc"],     #  the list of the sources to compile,
)

得到so 文件后,tf.load_op_library 动态加载so作为 module 使用(可以参考python module 动态加载加载)。

import tensorflow as tf
# 返回一个 A python module, containing the (op对应的)Python wrappers for Ops defined in the plugin.
# Python Module,是一个 Python 文件,以 .py 结尾,包含了 Python 对象定义和Python语句
zero_out_module = tf.load_op_library('zero_out.so')
with tf.Session():
  print(zero_out_module.zero_out([1,2,3,4,5])).eval() # eval 底层执行 session.run

zero_out_module.zero_out 可能仅用于演示,正规的ops 比如tf.matmul 等实现(对应到 编译tf生成的代码 gen_math_ops.py)会涉及到生成opDef (graphDef 的一部分)等逻辑。tensornet框架 自定义ops 示例

tensornet
    /core
        /kernels
            /sparse_table_ops.cc   # kernel实现
        /ops
            /sparse_table_ops.cc   # REGISTER_OP
        /BUILD                      
    /tensornet
        /core
            /gen_sparse_table_ops.py    # Bazel tf_gen_op_wrapper_py生成

BUILD 文件内容 tcnsorflow c++ op 生成 python调用接口

// 生成  生成op的lib,即so文件
cc_library(
    name = "sparse_table_ops_kernels",
    srcs = [
        "kernels/sparse_table_ops_dummy.cc",
        "ops/sparse_table_ops.cc",
    ],
    hdrs = [
        "//core/utility:semaphore",
    ],
    linkstatic = 1,
    deps = [
        "@org_tensorflow//tensorflow/core:framework",
        "@org_tensorflow//tensorflow/core:lib",
        "@org_tensorflow//tensorflow/core:protos_all_cc",
    ],
    alwayslink = 1,
)
// 生成python的接口,即gen_sparse_table_ops.py 文件
tf_gen_op_wrapper_py(
    name = "sparse_table_ops",
    deps = [":sparse_table_ops_kernels"],
    cc_linkopts = ['-lrt']
)

TensorFlow模型准实时更新上线的设计与实现

定制好 op 后,如何替换模型计算图中原生的 op 呢?TensorFlow 在模型保存时,会生成 meta_graph_def 文件,文件内容是采用类似 json 的格式描述计算图的结构关系。当加载此文件时,TensorFlow 会根据文件中描述的结构信息构建出计算图。可以修改模型保存的 meta_graph_def 文件,将其中的 op 替换为我们定制的 op,同时修改每个 node 的 input 和 output 关系,以修改 op 之间的依赖关系。PS: 当然这里说的替换原有的op

horovod

很多机器学习框架都会采用如下套路:shell脚本(可选),python端 和 C++端。

  1. Shell脚本是启动运行的入口,负责解析参数,确认并且调用训练程序;
  2. Python是用户的接口,引入了C++库,封装了API,负责运行时和底层C++交互;
  3. C++实现底层训练逻辑;

深度学习分布式训练框架horovod (2)一从使用者角度切入

引入库的作用是获取到 C++ 的函数,并且用 python 封装一下,这样就可以在 python 世界使用 C++代码了。比如下文,python 的 _allreduce 函数就会把功能转发给 C++,由 MPI_LIB.horovod_allreduce 完成。

def _allreduce(tensor, name=None, op=Sum, prescale_factor=1.0, postscale_factor=1.0,
               ignore_name_scope=False):
    if name is None and not _executing_eagerly():
        name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)
    return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
                                     prescale_factor=prescale_factor,
                                     postscale_factor=postscale_factor,
                                     ignore_name_scope=ignore_name_scope)
## 初始化时执行
def _load_library(name):
    """Loads a .so file containing the specified operators.
    """
    filename = resource_loader.get_path_to_datafile(name)
    library = load_library.load_op_library(filename)
    return library

# Check possible symbol not found error from tensorflow version mismatch
try:
    MPI_LIB = _load_library('mpi_lib' + get_ext_suffix())
except Exception as e:
    check_installed_version('tensorflow', tf.__version__, e)
    raise e
else:
    check_installed_version('tensorflow', tf.__version__)
打赏
文章版权声明:除非注明,否则均为彭超的博客原创文章,转载或复制请以超链接形式并注明出处。
相关推荐

发表评论:

◎欢迎参与讨论,请在这里发表您的看法、交流您的观点。

猜你喜欢