FastXML代码解读-done

FastXML代码解读-持续更新

FastXML是一种经典的极限多标签学习的方法，作者提供了完整的C++代码。先前编译成功并运行该代码，且添加了调试支持。本文将系统研究FastXML的代码。

FastXML的主文件有4个：

1. fastXML.h; 头文件，描述fastXML的节点（Node）类，GTree类，以及Param类。
2. fastXML.cpp；fastXML类的实现文件
3. fastXML_train.cpp; 训练过程，包含main函数
4. fastXML_predict.cpp; 预测过程，包含main函数

fastXML_train.cpp

先看最简单的，这个文件没有太多要讲的，主要包含main函数，以及超参数说明和传递过程。
超参列表：
-T 线程数，作者增h加了并发支持
-s Starting tree index，猜测是训练树所在的下标，默认为0，这个一般不改。
-t 用于集成的树的数目，文中默认为50。
-b 特征向量的偏置(bias)，如果没理解错的话，用于${\mathbf{w}}^T\mathbf{x}$中的bias项。
-c 支持向量机的权重系数。作者将向量$\mathbf{w}$看成是支持向量，这个系数就是原文中的$C_{\delta }$，不过作者加了个概率平衡的技巧。i.e., $C_{\delta }(\pm 1)=c/P(\pm 1)$。
-m 叶子节点最大样本数 MaxLeaf
-l 保留叶子节点概率值为top-l的标签。因为每个叶子节点都有一组样本，为了找到top-k的relevent的标签，作者就保留了一个最大值l。
关键代码如下：

int main(int argc, char* argv[])
{
    if(argc < 4)
        help();
    // 特征数据文件读取
    string ft_file = string(argv[1]);
    check_valid_filename(ft_file, true);
    SMatF* trn_X_Xf = new SMatF(ft_file);

    // 标签数据文件读取
    string lbl_file = string(argv[2]);
    check_valid_filename(lbl_file, true);
    SMatF* trn_X_Y = new SMatF(lbl_file);

    // 模型输出文件
    string model_folder = string(argv[3]);
    check_valid_foldername(model_folder);

    // 解析超参
    Param param = parse_param(argc-4,argv+4);
    // 输入数据的特征维度
    param.num_Xf = trn_X_Xf->nr;
    // 标签维度
    param.num_Y = trn_X_Y->nr;
    param.write(model_folder+"/param");

    if( param.quiet )
        loglvl = LOGLVL::QUIET;

    _float train_time;
    // 训练过程
    // 参数包括，训练特征，训练标签，超参列表，输出模型的文件夹，训练时间
    train_trees( trn_X_Xf, trn_X_Y, param, model_folder, train_time );
    cout << "training time: " << train_time/3600.0 << " hr" << endl;

    // 空间释放
    delete trn_X_Xf;
    delete trn_X_Y;
}

fastXML.h

这个头文件定义了三个类，分别是Param、Node、GTree，以及一系列的函数。

Param类定义了一系列的参数:

_int num_Xf; // 特征维度
_int num_Y; // 标签维度
_float log_loss_coeff; // C_\delta
_int max_leaf; // 叶子节点最大样本数
_int lbl_per_leaf; // label-probability pair数量，参考-l参数
_float bias; // w^Tx的偏置bias
_int num_thread; // 多线程支持
_int start_tree; 
_int num_tree; // 用于集成的树的数量
_bool quiet;

Node类定义了树中每一个节点的必要参数：

_bool is_leaf; // 是否是叶子节点
_int pos_child; // 正例孩子节点的下标 参考GTree.nodes
_int neg_child; // 负例孩子节点的下标 参考GTree.nodes
_int depth; // 深度
VecI X; // 存放在当前节点中训练样本的下标集合
VecIF w; // 存放每个特征的权重，w中的每一个元素是一个pair类型，int表示特征下标，float表示对应的权重。
VecIF leaf_dist; // 如果该节点是叶子节点，存放P

_float get_ram(); // 获取该节点的存储空间大小。

friend istream& operator>>(istream& fin, Node& node); // 从文件流中读取结点信息，这个文件流来源于任意一个.tree文件

GTree定义了一个树的所有结点，主要包含3个参数：

_int num_Xf; // 特征维度
_int num_Y; // 标签维度
vector<GNode*> nodes; // 结点集合

GTree( string model_dir, _int tree_no ); // 从一个.tree文件中构造一个GTree。

重点说一下.tree文件的内容，对任意一个i.tree文件，它存放了训练完成后的第i颗树的内容。

Number of Nodes # 第一行存放该树包含的节点数
// 循环Number Of Nodes次
IsLeaf # 存放是否是叶子节点 0 No，1 Yes
Pos Neg # 正例/负例孩子的下标
Depth # 当前节点的深度，从0开始
N i1 i2 ... iN # 样本数 ij为样本下标
M i1:j1 i2:j2 ... iM:jM # （如果不是叶子节点）非0权重的维度 il为对应的特征下标，jl为对应特征的权重值
M i1:j1 i2:j2 ... iM:jM # (如果是叶子节点) 将存入Node.leaf_dist（这个暂时不懂）

在fastXML.h文件里面还包含一些函数的声明：

// 训练单颗树, tree_no指定了树的下标
// 对应Algorithm 1的parallel for里面的部分
Tree* train_tree(SMatF* trn_ft_mat, SMatF* trn_lbl_mat, Param& param, _int tree_no);

// 训练所有的树，将每棵树的训练过程放入线程池
// 对应Algorithm 1的parallel for
void train_trees( SMatF* trn_X_Xf, SMatF* trn_X_Y, Param& param, string model_dir, _float& train_time );

SMatF* predict_tree(SMatF* tst_ft_mat, Tree* tree, Param& param);
SMatF* predict_trees( SMatF* tst_X_Xf, Param& param, string model_dir, _float& prediction_time, _float& model_size );

_bool optimize_ndcg( SMatF* X_Y, VecI& pos_or_neg );
void calc_leaf_prob( Node* node, SMatF* X_Y, Param& param );
_bool optimize_log_loss( SMatF* Xf_X, VecI& y, VecF& C, VecIF& sparse_w, Param& param );
void setup_thread_locals( _int num_X, _int num_Xf, _int num_Y );
pairII get_pos_neg_count(VecI& pos_or_neg);
void test_svm( VecI& X, SMatF* X_Xf, VecIF& w, VecF& values );

fastXML.cpp

fastXML.cpp实现了fastXML.h里面申明的函数。下面分别介绍

1. train_trees
   对应Algorithm 1里面的parallel for. 关键：支持并发。
   假设线程数为10，总计需要训练50棵树，那么每个线程负责串行地训练5棵树。

2. train_trees_thread
   每个线程负责串行地训练多棵树。
   关键：setup_thread_local函数负责设定线程级变量. 以及全局变量train_time的加锁。

{   // 当lock离开此域的时候，mtx会被释放。
    // https://en.cppreference.com/w/cpp/thread/lock_guard
    lock_guard<mutex> lock(mtx);
    *train_time += timer.toc();
}

C++关键字thread_local定义了线程级全局变量，比如

thread_local VecI countmap; // 在不同线程内是相互独立的副本。

每个线程的随机数发生器都不一样，保证了线程之间的随机数生成不会相互干扰。
每棵树的随机数种子被设定为树的id.

3. train_tree
   训练单棵树，对应Algorithm 1 parallel for里面的部分。
   关键：树中的节点增长是迭代式的，没有用到递归但也起到了文中Algorithm 1的GROW-NODE-RECURSIVE递归的作用。参考:
   shrink_data_matrices：从给定的样本下标集合中构建新的稀疏矩阵。
   SMat.shrink_mat:由所有训练样本构成的稀疏矩阵（也就是该对象）中，产生给定样本下标集合cols的稀疏矩阵s_mat。适用于节点的生成。
   SMat.active_dims:从cols(样本下标集合或标签下标集合)中 分解 实际的维度下标集合和对应维度的数量。
   calc_leaf_prob:计算叶子节点的概率P。保留前l个最大的标签-概率对。
   split_node: 节点分割。该节点的每个样本将会被划分到正域或者负域，corresponding to VecI pos_or_neg。

train_tree核心代码片段：

// nodes迭代地增长，这里并没有用递归也能起到递归的作用
for(_int i=0; i<nodes.size(); i++)
{    
    // 获取当前节点并试图分割，显然如果当前节点的样本数小于max_leaf，就不会分割了。
    Node* node = nodes[i];
    VecI& n_X = node->X; // 获取当前节点的所有样本
    SMatF* n_trn_Xf_X; // 当前节点样本构成的稀疏矩阵，转置
    SMatF* n_trn_X_Y; // 当前节点的标签构成的稀疏矩阵
    VecI n_Xf; // 每个样本的实际特征数
    VecI n_Y; // 每个样本的实际标签数

    // 将所有训练样本中下标属于n_X的的样本/标签重组为新的稀疏矩阵n_trn_Xf_X/n_trn_X_Y
    shrink_data_matrices( trn_X_Xf, trn_X_Y, n_X, n_trn_Xf_X, n_trn_X_Y, n_Xf, n_Y );

    // 对应Algorithm 1的GROW_NODE_RECURSIVE
    if(node->is_leaf)
    {
        // 计算叶子节点的P
        calc_leaf_prob( node, n_trn_X_Y, param );
    }
    else 
    {
        VecI pos_or_neg; // 当前节点的每个样本划分到正域还是负域
        // 划分节点. 
        // 输入：node, n_trn_Xf_X, n_trn_X_Y, param; 
        // 输出: pos_or_neg
        bool success = split_node( node, n_trn_Xf_X, n_trn_X_Y, pos_or_neg, param );

        if(success)
        {
            VecI pos_X, neg_X;
            for(_int j=0; j<n_X.size(); j++)
            {
                _int inst = n_X[j]; 
                if( pos_or_neg[j]==+1 )
                    pos_X.push_back(inst); 
                else
                    neg_X.push_back(inst);
            }

            Node* pos_node = new Node( pos_X, node->depth+1, param.max_leaf ); // 正子节点
            nodes.push_back(pos_node);
            node->pos_child = nodes.size()-1; // pos_node所在下标

            Node* neg_node = new Node( neg_X, node->depth+1, param.max_leaf ); // 负子节点
            nodes.push_back(neg_node);
            node->neg_child = nodes.size()-1; // neg_node所在下标
        }
        else 
        {
            node->is_leaf = true;
            i--; // back to当前节点，并计算P。
        }
    }
    
    postprocess_node( node, trn_X_Xf, trn_X_Y, n_X, n_Xf, n_Y );

    delete n_trn_Xf_X;
    delete n_trn_X_Y;
}

4. split_node
   不仅负责样本的划分，还负责学习当前节点的$\mathbf{w}$。对应文中的Algorithm 2.

但是split_node和原文的Algorithm 2.并不一致。在Algorithm 2中，作者试图先完全优化$\mathbf{r},\delta$，然后优化$\mathbf{w}$。也就是$\mathbf{r},\delta$这两个目标的优化次数一般比$\mathbf{w}$要多。
原文中对这3个优化目标的整个的优化是循环进行的。
但split_node的实际实现并没有整体的循环，而是先交替优化$\mathbf{r},\delta$，再优化$\mathbf{w}$。这就结束了。也就是说$\mathbf{w}$的优化只有最多一次（不懂）。

/**
 * @brief 分割当前节点的所有样本，并学习当前节点的权重向量
 * 
 * @param node 当前节点
 * @param Xf_X 当前节点的样本构成的Sparse Matrix
 * @param X_Y 当前节点的样本的标签构成的Sparse Matrix
 * @param pos_or_neg 文中的delta向量
 * @param param 
 * @return _bool 
 */
_bool split_node( Node* node, SMatF* Xf_X, SMatF* X_Y, VecI& pos_or_neg, Param& param )
{
    _int num_X = Xf_X->nr;
    pos_or_neg.resize( num_X );  
 
    // 对应Algorithm 2的第2行，随机地给delta_i分配-1或者1
    for( _int i=0; i<num_X; i++ )
    {
        _llint r = reng();

        if(r%2)
            pos_or_neg[i] = 1;
        else
            pos_or_neg[i] = -1;
    }
    
    // one run of ndcg optimization
    bool success;
    // 优化r和delta，主要就是delta
    success = optimize_ndcg( X_Y, pos_or_neg );

    if(!success)
        return false;

    // 优化目标5中的C_\delta,
    VecF C( num_X );
    // 算delta中正例和负例的个数
    // num_pos_neg.first表示pos，num_pos_neg.second表示neg
    pairII num_pos_neg = get_pos_neg_count( pos_or_neg );
    _float frac_pos = (_float)num_pos_neg.first/(num_pos_neg.first+num_pos_neg.second);
    _float frac_neg = (_float)num_pos_neg.second/(num_pos_neg.first+num_pos_neg.second);
    _double Cp = param.log_loss_coeff/frac_pos;
    _double Cn = param.log_loss_coeff/frac_neg;  // unequal Cp,Cn improves the balancing in some data sets
    
    // 根据正例负例的比例分配C_\delta，有点代价平衡的意思
    for( _int i=0; i<num_X; i++ )
        C[i] = pos_or_neg[i]==+1 ? Cp : Cn;

    // one run of log-loss optimization
    // 优化w
    success = optimize_log_loss( Xf_X, pos_or_neg, C, node->w, param );
    if(!success)
        return false;

    return true;
}

5. optimize_ndcg
   这个函数的功能是交替优化$r^{\pm }$和$\delta$。
   用到了两个关键的thread_local变量：1. discounts[l]存储$\frac{1}{1+l}$ for l in [1, L]; 2. csum_discount[l]存储discounts[1:l]的accumulation。局部变量idcg[i]存放当$||\mathbf{y}|{|}_1=i$时的$I_k(\mathbf{y})$

/**
 * @brief 优化r和delta
 * 
 * @param X_Y 样本标签
 * @param pos_or_neg delta
 * @return _bool 是否优化成功
 */
_bool optimize_ndcg( SMatF* X_Y, VecI& pos_or_neg )
{
    _int num_X = X_Y->nc; // 样本维度
    _int num_Y = X_Y->nr; // 标签维度
    _int* size = X_Y->size; 
    pairIF** data = X_Y->data; 
    
    _float eps = 1e-6;
    
    VecF idcgs( num_X );
    for( _int i=0; i<num_X; i++ )
        idcgs[i] = 1.0/csum_discounts[ size[i] ]; // I_k for each instance
    
    VecIF pos_sum( num_Y ); // r^+
    VecIF neg_sum( num_Y ); // r^-
    VecF diff_vec( num_Y );

    _float ndcg = -2;
    _float new_ndcg = -1;
    
    while(true)
    {
        // Step 1. 优化r^\pm
        for(_int i=0; i<num_Y; i++ )
        {
            pos_sum[i] = make_pair(i,0); 
            neg_sum[i] = make_pair(i,0);
            diff_vec[i] = 0;
        }

        for( _int i=0; i<num_X; i++ )
        {
            for( _int j=0; j<size[i]; j++ )
            {
                _int lbl = data[i][j].first; // lbl - 第i个样本的第j个标签下标
                _float val = data[i][j].second * idcgs[i]; // data[i][j].second为1

                if(pos_or_neg[i]==+1)  // 对应文中式11的求和
                    pos_sum[lbl].second += val;
                else
                    neg_sum[lbl].second += val;
            }
        }

        new_ndcg = 0;
        for(_int s=-1; s<=1; s+=2)
        {
            VecIF& sum = s==-1 ? neg_sum : pos_sum;
            // sort完成后，pos_sum和neg_sum里面的first序列存的就是$r^{\pm *}$
            sort(sum.begin(), sum.begin()+num_Y, comp_pair_by_second_desc<_int,_float>);

            for(_int i=0; i<num_Y; i++)
            {
                _int lbl = sum[i].first;
                _float val = sum[i].second;
                diff_vec[lbl] += s*discounts[i]; 
                new_ndcg += discounts[i]*val;
            }
        }
        new_ndcg /= num_X;

        // Step 2. 优化delta 
        for( _int i=0; i<num_X; i++ )
        {
            _float gain_diff = 0;
            for( _int j=0; j<size[i]; j++ )
            {
                _int lbl = data[i][j].first;
                _float val = data[i][j].second * idcgs[i];
                // 对应 (v_i^- - v_i^+)
                // 实际上v_i^\pm并没有考虑log item，而只考虑了nDCG item
                // 这是因为
                // 1. w被初始化为0, log item是只跟C_\delta(\pm 1)有关的常量
                // 2. 作者采用了balanced策略，使得C_\delta(\pm 1)相互抵消了。
                gain_diff += val*diff_vec[lbl]; 
            }

            if(gain_diff>0) // 对应文中式(13)
                pos_or_neg[i] = +1;
            else if(gain_diff<0)
                pos_or_neg[i] = -1;
        }
    
        if(new_ndcg-ndcg<eps)
            break; // 当nDCG变化不明显时，停止迭代
        else
            ndcg = new_ndcg;

    }

    pairII num_pos_neg = get_pos_neg_count(pos_or_neg);
    if(num_pos_neg.first==0 || num_pos_neg.second==0)
        return false;
    return true;
}

6. optimize_log_loss
   优化$\mathbf{w}$，这个暂时不懂，先码住。

7. predict_tree
   单棵树的预测过程，不需要训练集中的节点，比较简单。data[inst]存放第inst个测试样本在对应叶子节点的概率为top-l的label-property，按照下标排序，和leaf_dist一致。

for(_int i=0; i<nodes.size(); i++)
{
    Node* node = nodes[i];

    if(!node->is_leaf) // 迭代划分样本
    {
        VecI& X = node->X;
        // 线性分割
        test_svm(X, tst_X_Xf, node->w, values);
        for( _int j=0; j<X.size(); j++ )
            pos_or_neg[j] = values[j]>=0 ? +1 : -1;
        Node* pos_node = nodes[node->pos_child];
        pos_node->X.clear(); 
        Node* neg_node = nodes[node->neg_child];
        neg_node->X.clear(); 

        for(_int j=0; j<X.size(); j++)
        {
            if(pos_or_neg[j]==+1)
                pos_node->X.push_back(X[j]);
            else
                neg_node->X.push_back(X[j]);
        }
    }
    else // 到叶子节点时，计算每个样本的top-l标签的概率。
    {
        VecI& X = node->X;
        VecIF& leaf_dist = node->leaf_dist;
        _int* size = tst_score_mat->size;
        pairIF** data = tst_score_mat->data;

        for(_int j=0; j<X.size(); j++)
        {
            _int inst = X[j];
            size[inst] = leaf_dist.size();
            data[inst] = new pairIF[leaf_dist.size()];

            for(_int k=0; k<leaf_dist.size(); k++)
                data[inst][k] = leaf_dist[k];
        }
    }
}