如何在并行计算中直接写入sp.lil稀疏矩阵？

问题背景 #

我需要对大规模稀疏共现矩阵执行PPMI计算，当前实现通过中间列表存储所有并行计算结果，内存占用超过500GB，完全无法处理。核心需求是跳过中间列表存储，直接将PPMI计算结果写入稀疏矩阵，避免不必要的内存消耗。

共现矩阵定义：

cooc_matrix = sp.coo_matrix((data, (rows, cols)), shape=(num_vocab, num_vocab), dtype=np.float32)

当前低效实现（核心问题：中间列表ppmi_results占用巨量内存）：

from math import log
import numpy as np
import pandas as pd
import scipy.sparse as sp
import sys
from multiprocessing import Pool 
sys.path.append('/Users/myfunctions/')
from ppmi_calculation import compute_ppmi, ppmi_scorer 

if __name__ == '__main__':
    num_processes = 2
    pool = Pool(processes=num_processes)

    ppmi_results = []
    for i, j, v in zip(cooc_matrix.row, cooc_matrix.col, cooc_matrix.data):
        ppmi_results.append(pool.apply_async(compute_ppmi, args=(i, j, v, vocab_df, len_vocab, min_count, corpus_word_count)))

    pool.close()
    pool.join()

ppmi_matrix = sp.lil_matrix((len_vocab, len_vocab), dtype=np.float32)
for result in ppmi_results:
    i, j, ppmi_score = result.get()
    ppmi_matrix[i, j] = ppmi_score
    ppmi_matrix[j, i] = ppmi_score

配套的ppmi_calculation.py：

from math import log

ZERO = 0

def ppmi_scorer(worda_count, wordb_count, bigram_count, len_vocab, min_count, corpus_word_count):
    if bigram_count >= min_count:
        corpus_word_count = float(corpus_word_count)
        pa = worda_count / corpus_word_count
        pb = wordb_count / corpus_word_count
        pab = bigram_count / corpus_word_count
        try:
            return log(pab / (pa * pb)) / -log(pab)
        except ValueError:
            return ZERO
    else:
        return ZERO

def compute_ppmi(i, j, v, vocab_df, len_vocab, min_count, corpus_word_count):
    worda_count = vocab_df.iloc[i, 0]
    wordb_count = vocab_df.iloc[j, 0]
    ppmi_score = ppmi_scorer(worda_count, wordb_count, v, len_vocab, min_count, corpus_word_count)
    return i, j, ppmi_score

优化方案 #

核心思路：流式处理并行计算结果，仅保留非零PPMI条目，最后批量构建稀疏矩阵，彻底避免中间列表的内存占用。

步骤1：优化并行任务传递与计算 #

将vocab_df转换为numpy数组（减少进程间序列化开销），并使用imap_unordered流式获取计算结果，而非存储所有异步任务对象。

步骤2：实时过滤非零结果 #

主进程仅收集PPMI得分大于0的条目（稀疏矩阵无需存储0值），同时添加对称位置的条目（因为共现矩阵是对称的）。

步骤3：批量构建COO稀疏矩阵 #

COO格式是Scipy中最适合批量构建的稀疏矩阵类型，内存效率远高于逐点写入LIL矩阵。

完整优化代码 #

from math import log
import numpy as np
import pandas as pd
import scipy.sparse as sp
from multiprocessing import Pool 
sys.path.append('/Users/myfunctions/')
from ppmi_calculation import compute_ppmi, ppmi_scorer 

if __name__ == '__main__':
    num_processes = 2
    
    # 将词汇表计数转为numpy数组，减少进程间传递开销
    vocab_counts = vocab_df.iloc[:, 0].values
    
    # 准备并行任务参数迭代器
    task_args = [
        (i, j, v, vocab_counts, len_vocab, min_count, corpus_word_count)
        for i, j, v in zip(cooc_matrix.row, cooc_matrix.col, cooc_matrix.data)
    ]

    # 初始化存储非零PPMI结果的列表（仅保留有效数据）
    ppmi_rows = []
    ppmi_cols = []
    ppmi_data = []

    # 并行计算+流式处理结果
    with Pool(processes=num_processes) as pool:
        # 迭代获取每个任务的结果，实时处理
        for i, j, score in pool.imap_unordered(compute_ppmi, task_args):
            if score > 0:  # 过滤零值，利用稀疏特性压缩内存
                ppmi_rows.append(i)
                ppmi_cols.append(j)
                ppmi_data.append(score)
                # 添加对称位置的条目（共现矩阵对称）
                ppmi_rows.append(j)
                ppmi_cols.append(i)
                ppmi_data.append(score)

    # 批量构建COO格式PPMI矩阵（内存效率最高）
    ppmi_matrix = sp.coo_matrix(
        (ppmi_data, (ppmi_rows, ppmi_cols)),
        shape=(len_vocab, len_vocab),
        dtype=np.float32
    )

    # 按需转换为其他格式（如CSR，适合后续矩阵运算）
    ppmi_matrix = ppmi_matrix.tocsr()

配套修改compute_ppmi函数 #

def compute_ppmi(i, j, v, vocab_counts, len_vocab, min_count, corpus_word_count):
    # 直接从numpy数组取值，比DataFrame iloc更高效
    worda_count = vocab_counts[i]
    wordb_count = vocab_counts[j]
    ppmi_score = ppmi_scorer(worda_count, wordb_count, v, len_vocab, min_count, corpus_word_count)
    return i, j, ppmi_score

方案优势 #

1. 内存占用骤降：仅存储非零PPMI条目，避免了存储所有异步任务对象或计算结果，内存占用仅为原方案的极小部分。
2. 计算效率提升：使用numpy数组传递词汇计数，减少进程间序列化开销；imap_unordered流式处理无需等待所有任务完成。
3. 稀疏矩阵构建高效：COO格式批量构建比逐点写入LIL矩阵快数倍，且内存布局更紧凑。

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