如何加速计算Pandas DataFrame中文本的Word Mover's Distance？

核心问题原因 #

原代码慢主要有三个原因：

• 1000行数据两两计算总共有100万次WMD调用，WMD本身是线性规划类运算，复杂度高，原串行计算方式效率极低
• 重复计算：WMD具有对称性（A到B的距离等于B到A的距离），原代码全量计算了两遍所有配对，还有对角线自己和自己的0值也做了无效计算
• 冗余操作：每次循环都重复调用str.split()分词，且额外初始化Word2Vec模型完全没有必要，预训练的Glove向量可以直接调用WMD接口

优化后方案 #

第一步：前置预处理减少无效计算 #

提前完成所有句子的分词、无效词过滤，只保留预训练向量中存在的词，减少单次WMD的计算量。

第二步：利用对称性减少50%计算量 #

只计算上三角的配对，跳过对角线为0的结果，直接把总计算量从100万降到499500次。

第三步：多进程并行加速CPU密集计算 #

用多进程处理所有配对的计算，最终直接生成要求的DataFrame格式结果。

完整可运行代码 #

import pandas as pd
import numpy as np
import gensim.downloader as api
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm

# 全局加载预训练向量，避免多进程重复加载
wiki_vectors = api.load('glove-wiki-gigaword-50')

# 文本预处理函数，提前过滤无效词
def preprocess_sent(sent):
    # 可根据需求加小写化、去标点、去停用词等操作
    words = sent.strip().lower().split()
    return [w for w in words if w in wiki_vectors]

if __name__ == '__main__':
    # 读取你的原始数据
    data = pd.read_csv("你的数据文件路径.csv") # 替换为实际数据源
    # 提前预处理所有句子
    sent_list = [preprocess_sent(s) for s in data['Messages']]
    n = len(sent_list)
    # 生成需要计算的上三角配对，跳过对角线
    calc_pairs = [(i, j) for i in range(n) for j in range(i+1, n)]
    
    # 初始化距离矩阵，默认值为0（对角线自己和自己的距离就是0）
    dist_matrix = np.zeros((n, n))
    
    # 单配对WMD计算函数
    def calc_wmd(pair):
        i, j = pair
        return i, j, wiki_vectors.wmdistance(sent_list[i], sent_list[j])
    
    # 多进程计算，max_workers设置为你的CPU核心数即可
    with ProcessPoolExecutor(max_workers=8) as executor:
        # tqdm用来显示进度，不需要可以去掉，直接写list(executor.map(...))
        results = list(tqdm(executor.map(calc_wmd, calc_pairs), total=len(calc_pairs)))
    
    # 填充结果到矩阵，利用对称性同时填上下三角
    for i, j, dist in results:
        dist_matrix[i][j] = dist
        dist_matrix[j][i] = dist
    
    # 转换为带ID索引的DataFrame
    distance_results = pd.DataFrame(dist_matrix, index=data['ID'], columns=data['ID'])

额外可选优化 #

如果后续要处理更大规模的数据，可以做如下调整：

• 替换gensim自带的WMD为pot（Python Optimal Transport）库的优化实现，开启sinkhorn近似计算，速度可以再提升3-10倍，误差可控
• 提前对文本做停用词过滤、标点移除，进一步减少每个句子的有效词数量，降低单次WMD的计算复杂度
• 若数据量超过5000行，可考虑先对句子做聚类分组，只计算组内的WMD距离，避免全量配对的指数级计算量增长

内容的提问来源于stack exchange，提问作者mjoy