加载 lightgbm 模型并使用 predict 与并行 for 循环冻结（Python）

Question

我需要使用我的模型在 python 中批量和并行地进行预测。如果我加载模型并在常规 for 循环中创建数据框并使用 predict 函数，它就没有问题。如果我在 python 中使用 multiprocessing 并行创建不相交的数据帧，然后使用 predict 函数 for 循环无限期冻结。为什么会出现这种行为？

这是我的代码片段：

with open('models/model_test.pkl', 'rb') as fin:
    pkl_bst = pickle.load(fin)

def predict_generator(X):

    df = X

    print(df.head())
    df = (df.groupby(['user_id']).recommender_items.apply(flat_map)
          .reset_index().drop('level_1', axis=1))
    df.columns = ['user_id', 'product_id']

    print('Merge Data')
    user_lookup = pd.read_csv('data/user_lookup.csv')
    product_lookup = pd.read_csv('data/product_lookup.csv')
    product_map = dict(zip(product_lookup.product_id, product_lookup.name))

    print(user_lookup.head())

    df = pd.merge(df, user_lookup, on=['user_id'])
    df = pd.merge(df, product_lookup, on=['product_id'])
    df = df.sort_values(['user_id', 'product_id'])

    users = df.user_id.values
    items = df.product_id.values
    df.drop(['user_id', 'product_id'], axis=1, inplace=True)

    print('Prediction Step')

    prediction = pkl_bst.predict(df, num_iteration=pkl_bst.best_iteration)
    print('Prediction Complete')

    validation = pd.DataFrame(zip(users, items, prediction),
                              columns=['user', 'item', 'prediction'])
    validation['name'] = (validation.item
                          .apply(lambda x: get_mapping(x, product_map)))
    validation = pd.DataFrame(zip(validation.user,
                              zip(validation.name,
                                  validation.prediction)),
                              columns=['user', 'prediction'])
    print(validation.head())

    def get_items(x):

        sorted_list = sorted(list(x), key=lambda i: i[1], reverse=True)[:20]
        sorted_list = random.sample(sorted_list, 10)
        return [k for k, _ in sorted_list]

    relevance = validation.groupby('user').prediction.apply(get_items)
    return relevance.reset_index()

这有效但非常慢：

results = []
for d in df_list_sub:
    r = predict_generator(d)
    results.append(r)

这打破了：

from multiprocessing import Pool
import tqdm
pool = Pool(processes=8)
results = []
for x in tqdm.tqdm(pool.imap_unordered(predict_generator, df_list_sub), total=len(df_list_sub)):
    results.append(x)
    pass
pool.close()
pool.join()

如果有人可以帮助我，我将不胜感激。

Answer 1

我自己也偶然发现了这一点。这是因为 LightGBM 只允许从单个进程访问预测函数。开发人员明确添加了此逻辑，因为从多个进程调用预测函数没有意义，因为预测函数已经利用了所有可用的 CPU。除此之外，允许多进程预测可能会导致性能更差。有关此内容的更多信息可以在此 GitHub 问题中找到。