tensornet源码调试解析

最近阅读了tensornet的源码，其设计思想很值得借鉴。对于架构设计的感兴趣的同学，强烈建议阅读一下。tensornet在tensorflow的基础上进行二次开发，针对广告推荐等大规模稀疏场景优化的分布式训练框架。相比之前的分布式框架，tensornet借助mpi集群管理，每个节点单独维护一个ps，省去了维护管理节点的成本。本文主要记录了我个人对tensornet源码的阅读以及理。

1. 环境准备

1.1 安装

大家可以根据tensornet提供的dockerfile创建一个tensornet的镜像，我本人使用的是vscode编辑器，这个和docker兼容做的非常棒，不需要像pycharm等配置ssh连接。另一个最大的好处是支持python和c++混合调试的方式。创建镜像，vscode以及相应的插件安装(我安装了vscode的docker插件)，大家可以百度。安装好后可以在vscode中打开，就可以像在本地一样，调试编辑代码了。同时还需要安装linux 的gdb-debug 工具，python debug工具，这些同样都可以百度到。

1.2 调试

为了debug断点调试代码，需要做以下两件事：

1.2.1 创建调试入口文件

在根目录下创建 .vscode 文件夹，在里面创建launch.json文件,文件内容如下：

{
    "version": "0.2.0",
    "configurations": [
        {
            "name": "Python: Current File",
            "type": "python",
            "request": "launch",
            "program": "/tensornet/debug2.py",
            "console": "integratedTerminal",
            "port": 10010,
            "host": "localhost",
            "pathMappings": [
                {
                    "localRoot": "${workspaceFolder}", // You may also manually specify the directory containing your source code.
                    "remoteRoot": "${workspaceFolder}", // Linux example; adjust as necessary for your OS and situation.
                }
            ]
        },
        {
            "name": "(gdb) Attach",
            "type": "cppdbg",
            "request": "attach",
            "program": "/usr/bin/python3.7",
            "processId": "${command:pickProcess}",
            "MIMode": "gdb",
            "sourceFileMap":{
                "/proc/self/cwd/": "/tensornet"
            },
            "setupCommands": [
                {
                    "description": "Enable pretty-printing for gdb",
                    "text": "-enable-pretty-printing",
                    "ignoreFailures": true
                }
            ]
        }
    ]
}

debug2.py 是python调试入口文件，可以将官方的例子 01-begin-with-wide-deep.ipynb 内容拷贝这个python文件

1.2.2 编译调试库

直接按照官方给的编译命令进行编译，是无法调试c++代码的，里面不包含调试信息。执行如下命令，编译调试动态库。

bazel build -c dbg --sandbox_debug --copt -g //core:_pywrap_tn.so --verbose_failures
cp -f bazel-bin/core/_pywrap_tn.so tensornet/core 

1.2.3 测试debug

有了上面的准备工作，就可以测试python 和c++混合debug了。python只是提供了接口层，python运行主要是为了构建tensorflow 的图，真实的运算逻辑都在底层的c++中。我们以整个系统的初始化为例。py_wrapper.cc 提供了python调用底层c++的接口，在此处打上断点：
ps_strategy.py 是接口层的入口，通过tn.core.init()调用底层的c++对整个集群进行初始化。

下面给出一个完整debug的例子

tensornet断点调试

2. 源码阅读

tensornet在tensorflow基础上进行二次开发，代码量不大，没有tensorflow定制化开发经验的情况下，理解起来也很困难。但是再复杂的程序，肯定也有入口函数，我们只需要追踪其入口函数，即可理清其逻辑。

官方给的wide-deep示例中，首先会引入import tensornet as tn，这步操作会执行tensornet下面的_init_.py 文件，这个属于python的基础语法。在步操作中，只是简单的将封装的python api对外进行暴露，没有其他实际的操作。那些函数定义，可以暂时不看，因为没有调用，也就没有执行。真正开始执行的是 train(tn.distribute.PsStrategy()) ，关键点在tn.distribute.PsStrategy()。找到对应的函数就会眼前一亮，tn.core.init()，在这里执行了整个集群的初始化。

2.1 集群初始化

PsStrategy 继承了OneDeviceStrategy，这个是tensorflow分布式训练策略接口，可能会底层一些。tensornet继承重写了自己的分布式策略，核心就在这个 tn.core.init() 化方法中。 这个方法是python的调用方式，追踪到tensornet的core包下，只有一个.so文件和一些脚本自动生成的python文件，这些自动生成的python文件是利用tf_gen_op_wrapper_py根据c++文件自动封装生成的接口文件，可以暂时不用管。那么这个方法就一定是调用的_pywrap_tn.so文件。那么就一定有生成这个so文件对应的源码，翻看c++的core文件，就会发现如下的文件：
此时就会恍然大会，这个就是将c++文件封装成供python调用的so文件。主要使用的pybind11这个工具，大家百度查看这个工具的几个demo，就会明白其用法。看完这个工具的用法，你就会明白为啥生成的so文件叫_pywrap_tn。在这个文件第一个封装的方法就是init函数，此时线索就很明显了。我们可以看一下在这个init方法中，都做了什么操作。

		// 获取ps集群的实例
        PsCluster* cluster = PsCluster::Instance();

        //判断集群是否初始化
        if (cluster->IsInitialized()) {
            return true;
        }
        
        //对ps集群进行初始化操作
        if (cluster->Init() < 0) {
            throw py::value_error("Init tensornet fail");
        }

        //BalanceInputDataInfo的初始化
        tensorflow::BalanceInputDataInfo* data_info = tensorflow::BalanceInputDataInfo::Instance();

        if (data_info->Init(cluster->RankNum(), cluster->Rank()) < 0) {
            throw py::value_error("Init BalanceInputDataInfo fail");
        }

        return true;

首先会获取集群的单例对象，然后调用集群的初始化方法。在初始化方法中，绑定brpc的服务–server_->AddService(&ps_service_impl_, brpc::SERVER_DOESNT_OWN_SERVICE)。然后就是获取mpi集群管理器，tensornet是利用mpi进行集群管理的。先将mpi集群初始化，这个具体里面的一些细节(节点的数量，每个节点的端口，ip等信息)大家可以自己去查阅。集群初始化完毕，会通过InitRemoteServers_()建立节点之间的相互通讯链接，
这些链接都是rpc通讯。总结来说，在初始化方法中，主要做了以下三件事情：

• 启动本机rpc服务
• 初始化mpi集群管理器
• 建立节点之前的相互通讯

2.2 CategoryColumn封装

接下来执行的是columns_builder，在该方法中，会根据配置生成tn.feature_column.category_column 生成category_column对象。在CategoryColumn类中，继承了父类的一些方法。目前为止还看不出来这些方法具体的用途，可以暂时跳过。create_model中，会调用tn.core.AdaGrad方法，生成优化器。这个又是c++类的接口，同样可以在py_wrapper.cc中找到对其封装。在封装的方法中，会获取python传入的参数，转换成c++类型，调用C++ 的AdaGrad类。将该对象的指针引用返回给python接口。

2.3 EmbeddingFeatures封装

在EmbeddingFeatures类中，传入feature_columns和优化器。在初始化方法中，会检查传入的feature_columns embedding维度是否一致。目前tensornet只能支持维度一致的id词表，接下来会创建最重要的一个对象–StateManagerImpl。

在这StateManagerImpl的构造函数中，又会调用c++的接口create_sparse_table，创建稀疏表。在构造函数中，除了一系列的初始化，最主要的操作是创建id表的存储空间SparseOptimizerKernel。默认是创建八个块进行存储，每个块是一个SparseKernelBlock对象。在每个块中利用一个map存储的embedding矩阵，key是id，value是一个封装的SparseAdaGradValue结构体。里面存储了当前embedding的维度，以及优化器的参数。具体的细节可以跟踪源码阅读，这个方法最后会返回一个表的句柄，供以后查找使用。

创建完_state_manager对象后，会调用EmbeddingFeatures的build方法，在build方法中，会依次调用每一列的create_state方法，这个方法中会调用_state_manager的create_variable方法，会生成词表，这个词表只是起到一个占位符的作用，实际更新的参数在ps上，每次拉取更新的时候，会将当前batch中出现的词表拷贝到这个词表中，会根据索引在这个变量中查找。会将创建的词表var的索引保存到字典中，以供后续查找。

下一步就是调用EmbeddingFeatures的call方法，首先会过滤不是id特征的列，然后将使用的id 特征调用FeatureTransformationCache缓存起来。最重要的一个方法self._state_manager.pull ,在这个方法中，会调用自定义的op，拉取稀疏特征对应的embedding向量。如何自定义tensorflow op，可以查阅tensorflow官方文档。拉取id特征的embedding向量的op在sparse_table_ops.cc中，sparse_table_ops_dummy.cc只是起到了编译op对应的Python文件接口使用，实际运行中使用的逻辑在sparse_table_ops文件中。具体细节就不在这里展开了，主要的操作如下：

1. 将id(整数)依次给个序号，相同id序号一样，这个序号代表当前embedding向量在create_variable创建的var对象中的索引。
2. 根据id hash将id均匀分到当前节点所有的机器上，查询本地的词表或者利用rpc查询其他节点
3. 第一次查询的时候的时候 会分配对应的内存，并对内存初始化
4. 将多个节点查询的向量，拷贝到var的内存中

拉取向量后，会调用get_dense_tensor方法，将稀疏的embedding向量，转化为dense embedding。考虑到有可能有多个id，所以这一步转换目的主要是变成一个embedding向量，可以按照mean、sum等方式。在这里前面定义的那些方法都会被调用到。

backwards方法中，主要是反向传播的时候用到，更新id embedding向量。这个也同样可以追踪相关的c++文件。

2.4 优化器封装

在优化器optimizer.py文件中，主要是为了更新dense参数(全连接网络的参数)。
在优化器方法中，最重要的一个方法是_distributed_apply，这个方法会调用_resource_apply_dense和_resource_apply_sparse，因为tensornet参数都在ps上，所以这两个方法没有具体的逻辑。因此必须要重写父类的两个方法，不然就会调用者两个方法的逻辑。_distributed_apply会把dense梯度传到服务器上进行更新。

2.5 模型的封装

在Model.py中，最重要的一个方法是train_step ,这个是训练每一步的逻辑。 self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))这个只会调用优化器更新dense参数，然后调用backwards方法更新id特征的参数，具体的更新方法在EmbeddingFeatures中。