Python中基于gensim Word2Vec的多进程批量任务突然中止求助

嘿，我碰到过好几次Gensim WMD在多进程环境里突然挂掉的情况，大概率是模型实例跨进程传递或者资源共享的坑导致的，给你几个靠谱的解决方案：

核心原因分析 #

Gensim的KeyedVectors和WmdSimilarity实例不能直接跨进程传递：

• 如果你在主进程加载模型、创建WmdSimilarity实例，再把它传给子进程，会因为底层内存映射资源（比如GoogleNews模型的内存映射文件）在fork/spawn进程时失效，或者pickle序列化失败导致崩溃。
• 另外，WMD计算本身比较耗内存和CPU，批量处理时如果某个查询触发未捕获的异常（比如预处理错误、内存溢出），也会导致进程直接中止。

解决方案1：每个进程独立初始化模型和WMD实例 #

最稳妥的方式是让每个子进程自己加载模型、创建WmdSimilarity实例，避免跨进程共享实例。可以用进程池的initializer参数在进程启动时完成初始化：

import multiprocessing
from gensim.models import KeyedVectors
from gensim.similarities import WmdSimilarity

# 子进程全局变量，每个进程会初始化自己的副本
model = None
wmd_instance = None

def init_worker(processed_set):
    """每个子进程启动时初始化模型和WMD实例"""
    global model, wmd_instance
    # 每个进程独立加载模型（注意：GoogleNews模型很大，内存要够）
    model = KeyedVectors.load_word2vec_format(
        '/users/myuser/method_approaches/google_news_requirements/GoogleNews-vectors-negative300.bin.gz',
        binary=True
    )
    wmd_instance = WmdSimilarity(processed_set, model, num_best=10)

def process_single_query(query):
    """子进程中处理单个查询的函数"""
    try:
        # 这里替换成你的查询预处理逻辑
        processed_query = your_preprocess_function(query)
        # 计算相似性
        return wmd_instance[processed_query]
    except Exception as e:
        # 捕获所有异常，避免进程直接崩溃
        import traceback
        print(f"处理查询时出错：{query}")
        traceback.print_exc()
        return None

if __name__ == '__main__':
    # 假设这是你的预处理后语料库
    processed_set = [...]
    # 你的批量查询列表
    query_list = [...]
    
    # 创建进程池，每个进程初始化模型和WMD实例
    # 注意：进程数不要开太多，GoogleNews模型每个进程占3-4G内存
    with multiprocessing.Pool(
        processes=2,  # 根据你的内存调整进程数
        initializer=init_worker,
        initargs=(processed_set,)
    ) as pool:
        # 批量处理查询
        results = pool.map(process_single_query, query_list)
    
    # 处理返回结果
    for idx, res in enumerate(results):
        if res is not None:
            print(f"查询{idx}的相似结果：{res}")

解决方案2：优化内存占用（可选） #

如果你的机器内存不够，每个进程加载GoogleNews模型太吃力，可以试试这些优化：

• 用limit参数加载部分词向量：比如KeyedVectors.load_word2vec_format(..., limit=500000)只加载前50万个高频词，大幅减少内存占用。
• 换成更小的预训练模型，比如GloVe的6B词向量（约1.5G），或者对应语言的轻量预训练模型。

解决方案3：捕获所有异常定位问题 #

如果进程还是突然中止，大概率是某个查询触发了未捕获的异常。一定要在处理函数里加上try-except块，打印详细的错误栈，这样能快速定位问题（比如预处理逻辑有bug、查询包含模型中没有的词导致崩溃等）。

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