lightgbm用於排序

一. 

　　LTR(learning to rank)經常用於搜索排序中，開源工具中比較有名的是微軟的ranklib，但是這個好像是單機版的，也有好長時間沒有更新了。所以打算想利用lightgbm進行排序，但網上關於lightgbm用於排序的代碼很少，關於回歸和分類的倒是一堆。這里我將貼上python版的lightgbm用於排序的代碼，里面將包括訓練、獲取葉結點、ndcg評估、預測以及特征重要度等處理代碼，有需要的朋友可以參考一下或進行修改。

　　其實在使用時，本人也對比了ranklib中的lambdamart和lightgbm，令人映像最深刻的是lightgbm的訓練速度非常快，快的起飛。可能lambdamart訓練需要幾個小時，而lightgbm只需要幾分鍾，但是后面的ndcg測試都差不多，不像論文中所說的lightgbm精度高一點。lightgbm的訓練速度快，我想可能最大的原因要可能是：a.節點分裂用到了直方圖，而不是預排序方法；b.基於梯度的單邊采樣，即行采樣；c.互斥特征綁定，即列采樣；d.其於leaf-wise決策樹生長策略；e.類別特征的支持等

二.代碼

第一部分代碼塊是主代碼，后面三個代碼塊是用到的加載數據和ndcg。運行主代碼使用命令如訓練模型使用：python lgb.py -train等

完成代碼和數據格式放在https://github.com/jiangnanboy/learning_to_rank上面，大家可以參考一下！！！！！

import os
import lightgbm as lgb
from sklearn import datasets as ds
import pandas as pd

import numpy as np
from datetime import datetime
import sys
from sklearn.preprocessing import OneHotEncoder

def split_data_from_keyword(data_read, data_group, data_feats):
    '''
    利用pandas
    轉為lightgbm需要的格式進行保存
    :param data_read:
    :param data_save:
    :return:
    '''
    with open(data_group, 'w', encoding='utf-8') as group_path:
        with open(data_feats, 'w', encoding='utf-8') as feats_path:
            dataframe = pd.read_csv(data_read,
                                    sep=' ',
                                    header=None,
                                    encoding="utf-8",
                                    engine='python')
            current_keyword = ''
            current_data = []
            group_size = 0
            for _, row in dataframe.iterrows():
                feats_line = [str(row[0])]
                for i in range(2, len(dataframe.columns) - 1):
                    feats_line.append(str(row[i]))
                if current_keyword == '':
                    current_keyword = row[1]
                if row[1] == current_keyword:
                    current_data.append(feats_line)
                    group_size += 1
                else:
                    for line in current_data:
                        feats_path.write(' '.join(line))
                        feats_path.write('\n')
                    group_path.write(str(group_size) + '\n')

                    group_size = 1
                    current_data = []
                    current_keyword = row[1]
                    current_data.append(feats_line)

            for line in current_data:
                feats_path.write(' '.join(line))
                feats_path.write('\n')
            group_path.write(str(group_size) + '\n')

def save_data(group_data, output_feature, output_group):
    '''
    group與features分別進行保存
    :param group_data:
    :param output_feature:
    :param output_group:
    :return:
    '''
    if len(group_data) == 0:
        return
    output_group.write(str(len(group_data)) + '\n')
    for data in group_data:
        # 只包含非零特征
        # feats = [p for p in data[2:] if float(p.split(":")[1]) != 0.0]
        feats = [p for p in data[2:]]
        output_feature.write(data[0] + ' ' + ' '.join(feats) + '\n') # data[0] => level ; data[2:] => feats

def process_data_format(test_path, test_feats, test_group):
    '''
     轉為lightgbm需要的格式進行保存
     '''
    with open(test_path, 'r', encoding='utf-8') as fi:
        with open(test_feats, 'w', encoding='utf-8') as output_feature:
            with open(test_group, 'w', encoding='utf-8') as output_group:
                group_data = []
                group = ''
                for line in fi:
                    if not line:
                        break
                    if '#' in line:
                        line = line[:line.index('#')]
                    splits = line.strip().split()
                    if splits[1] != group: # qid => splits[1]
                        save_data(group_data, output_feature, output_group)
                        group_data = []
                    group = splits[1]
                    group_data.append(splits)
                save_data(group_data, output_feature, output_group)

def load_data(feats, group):
    '''
    加載數據
    分別加載feature,label,query
    '''
    x_train, y_train = ds.load_svmlight_file(feats)
    q_train = np.loadtxt(group)
    return x_train, y_train, q_train

def load_data_from_raw(raw_data):
    with open(raw_data, 'r', encoding='utf-8') as testfile:
        test_X, test_y, test_qids, comments = letor.read_dataset(testfile)
    return test_X, test_y, test_qids, comments

def train(x_train, y_train, q_train, model_save_path):
    '''
    模型的訓練和保存
    '''
    train_data = lgb.Dataset(x_train, label=y_train, group=q_train)
    params = {
        'task': 'train',  # 執行的任務類型
        'boosting_type': 'gbrt',  # 基學習器
        'objective': 'lambdarank',  # 排序任務(目標函數)
        'metric': 'ndcg',  # 度量的指標(評估函數)
        'max_position': 10,  # @NDCG 位置優化
        'metric_freq': 1,  # 每隔多少次輸出一次度量結果
        'train_metric': True,  # 訓練時就輸出度量結果
        'ndcg_at': [10],
        'max_bin': 255,  # 一個整數，表示最大的桶的數量。默認值為 255。lightgbm 會根據它來自動壓縮內存。如max_bin=255 時，則lightgbm 將使用uint8 來表示特征的每一個值。
        'num_iterations': 500,  # 迭代次數
        'learning_rate': 0.01,  # 學習率
        'num_leaves': 31,  # 葉子數
        # 'max_depth':6,
        'tree_learner': 'serial',  # 用於並行學習，‘serial’： 單台機器的tree learner
        'min_data_in_leaf': 30,  # 一個葉子節點上包含的最少樣本數量
        'verbose': 2  # 顯示訓練時的信息
    }
    gbm = lgb.train(params, train_data, valid_sets=[train_data])
    gbm.save_model(model_save_path)

def predict(x_test, comments, model_input_path):
    '''
    預測得分並排序
    '''
    gbm = lgb.Booster(model_file=model_input_path)  # 加載model

    ypred = gbm.predict(x_test)

    predicted_sorted_indexes = np.argsort(ypred)[::-1]  # 返回從大到小的索引

    t_results = comments[predicted_sorted_indexes]  # 返回對應的comments,從大到小的排序

    return t_results

def test_data_ndcg(model_path, test_path):
    '''
    評估測試數據的ndcg
    '''
    with open(test_path, 'r', encoding='utf-8') as testfile:
        test_X, test_y, test_qids, comments = letor.read_dataset(testfile)

    gbm = lgb.Booster(model_file=model_path)
    test_predict = gbm.predict(test_X)

    average_ndcg, _ = ndcg.validate(test_qids, test_y, test_predict, 60)
    # 所有qid的平均ndcg
    print("all qid average ndcg: ", average_ndcg)
    print("job done!")

def plot_print_feature_importance(model_path):
    '''
    打印特征的重要度
    '''
    #模型中的特征是Column_數字,這里打印重要度時可以映射到真實的特征名
    feats_dict = {
        'Column_0': '特征0名稱',
        'Column_1': '特征1名稱',
        'Column_2': '特征2名稱',
        'Column_3': '特征3名稱',
        'Column_4': '特征4名稱',
        'Column_5': '特征5名稱',
        'Column_6': '特征6名稱',
        'Column_7': '特征7名稱',
        'Column_8': '特征8名稱',
        'Column_9': '特征9名稱',
        'Column_10': '特征10名稱',
    }
    if not os.path.exists(model_path):
        print("file no exists! {}".format(model_path))
        sys.exit(0)

    gbm = lgb.Booster(model_file=model_path)

    # 打印和保存特征重要度
    importances = gbm.feature_importance(importance_type='split')
    feature_names = gbm.feature_name()

    sum = 0.
    for value in importances:
        sum += value

    for feature_name, importance in zip(feature_names, importances):
        if importance != 0:
            feat_id = int(feature_name.split('_')[1]) + 1
            print('{} : {} : {} : {}'.format(feat_id, feats_dict[feature_name], importance, importance / sum))

def get_leaf_index(data, model_path):
    '''
    得到葉結點並進行one-hot編碼
    '''
    gbm = lgb.Booster(model_file=model_path)
    ypred = gbm.predict(data, pred_leaf=True)

    one_hot_encoder = OneHotEncoder()
    x_one_hot = one_hot_encoder.fit_transform(ypred)
    print(x_one_hot.toarray()[0])

if __name__ == '__main__':
    model_path = "保存模型的路徑"

    if len(sys.argv) != 2:
        print("Usage: python main.py [-process | -train | -predict | -ndcg | -feature | -leaf]")
        sys.exit(0)

    if sys.argv[1] == '-process':
        # 訓練樣本的格式與ranklib中的訓練樣本是一樣的,但是這里需要處理成lightgbm中排序所需的格式
        # lightgbm中是將樣本特征和group分開保存為txt的,什么意思呢,看下面解釋
        '''
        feats:
        1 1:0.2 2:0.4 ...
        2 1:0.2 2:0.4 ...
        1 1:0.2 2:0.4 ...
        3 1:0.2 2:0.4 ...
        group:
        2
        4
        這里group中2表示前2個是一個qid,4表示后兩個是一個qid
        '''
        raw_data_path = '訓練樣本集路徑'
        data_feats = '特征保存路徑'
        data_group = 'group保存路徑'
        process_data_format(raw_data_path, data_feats, data_group)

    elif sys.argv[1] == '-train':
        # train
        train_start = datetime.now()
        data_feats = '特征保存路徑'
        data_group = 'group保存路徑'
        x_train, y_train, q_train = load_data(data_feats, data_group)
        train(x_train, y_train, q_train, model_path)
        train_end = datetime.now()
        consume_time = (train_end - train_start).seconds
        print("consume time : {}".format(consume_time))

    elif sys.argv[1] == '-predict':
        train_start = datetime.now()
        raw_data_path = '需要預測的數據路徑'#格式如ranklib中的數據格式
        test_X, test_y, test_qids, comments = load_data_from_raw(raw_data_path)
        t_results = predict(test_X, comments, model_path)
        train_end = datetime.now()
        consume_time = (train_end - train_start).seconds
        print("consume time : {}".format(consume_time))

    elif sys.argv[1] == '-ndcg':
        # ndcg
        test_path = '測試的數據路徑'#評估測試數據的平均ndcg
        test_data_ndcg(model_path, test_path)

    elif sys.argv[1] == '-feature':
        plot_print_feature_importance(model_path)

    elif sys.argv[1] == '-leaf':
        #利用模型得到樣本葉結點的one-hot表示
        raw_data = '測試數據路徑'#
        with open(raw_data, 'r', encoding='utf-8') as testfile:
            test_X, test_y, test_qids, comments = letor.read_dataset(testfile)
        get_leaf_index(test_X, model_path)