多进程处理二元模型计数子矩阵为何出现异常输出？

问题背景 #

我写了一个训练二元模型的程序，为提升运行速度，拆分计数矩阵后用多进程处理子矩阵，再合并成概率矩阵。单进程版本运行正常，但多进程模式下，只有当n_processes = 1时结果正确；设置为2、8或16进程时，第一个子矩阵的计算是对的（每行概率和为1），其余子矩阵的数值完全混乱，范围在0.001到250.0之间，而且进程数增加也不会持续恶化。

子矩阵处理函数 #

def kn_probs(bigram_counts, type_counts, d=0.75):
"""
Calculate Kneser-Ney probabilities on a matrix of bigram counts.

:param bigram_counts: 2D matrix of integer counts of bigrams
:param type_counts: number of bigram types, list of row types
counts, list of col types counts.
:param d: discount value subtracted from the count of each bigram
:return: a 2D matrix containing the Kneser-Ney bigram probabilities
"""
n_bigram_types, row_types_counts, col_types_counts = type_counts

# initialise matrix
kn_probs = np.zeros((len(bigram_counts), len(bigram_counts[0])), dtype=np.float16)

# iterate through the matrix
for row_i, row in enumerate(bigram_counts):
    
    # compute the pre_word (i.e. row) terms
    pre_word_count = np.sum(row)
    n_pre_word_followers = row_types_counts[row_i]
    kn_lambda = (d / pre_word_count) * n_pre_word_followers
    
    # iterate through the row
    for col_i, bigram_count in enumerate(row):
        
        n_word_precedents = col_types_counts[col_i]
        
        # compute KN prob and save it
        kn_prob = (max(bigram_count-d, 0) / pre_word_count)\
                    + (kn_lambda * (n_word_precedents / n_bigram_types))
        kn_probs[row_i][col_i] = kn_prob

return kn_probs

多进程处理代码 #

# split bigram_counts to allow for multiprocessing
n_processes = 16
bigram_counts_chunks = np.array_split(bigram_counts, n_processes)
length = len(bigram_counts)
del(bigram_counts)

# create a pool of worker processes
with mp.Pool(n_processes) as p:
    
    # compute type counts separately and merge them (needed for Kneser-Ney smoothing)
    chunks_type_counts = p.map(count_types, bigram_counts_chunks)
    n_bigram_types = sum([type_counts[0] for type_counts in chunks_type_counts])
    row_types_counts = np.concatenate([type_counts[1] for type_counts in chunks_type_counts])
    col_types_counts = [0] * length
    for type_counts in chunks_type_counts:
        col_types_counts = np.add(col_types_counts, type_counts[2])
    type_counts = (n_bigram_types, row_types_counts, col_types_counts)
    
    # map the kn_probs() function to each chunk in parallel
    bigram_probs_chunks = p.starmap(kn_probs, zip(bigram_counts_chunks, itertools.repeat(type_counts)))
    del(bigram_counts_chunks)
    
# concatenate chunks into one array
bigram_probs = np.concatenate(bigram_probs_chunks)

环境：HP笔记本，Windows 11 64位，Python 3.11.0，16GB内存

问题根源 #

核心问题是行类型统计的索引不匹配：

• 拆分后的每个子矩阵，内部行索引是从0开始的局部索引，但kn_probs函数里用row_types_counts[row_i]时，row_types_counts是全局拼接后的数组，对应原矩阵的全局行索引。
• 第一个子矩阵的局部索引0_{k刚好和全局索引0}k重合，所以计算正确；但第二个子矩阵的局部索引0对应的是全局索引k+1，此时用局部索引取全局数组，拿到的是完全错误的n_pre_word_followers值，导致kn_lambda计算异常，最终概率值彻底混乱。

修复方案 #

给每个子矩阵传递对应的局部行类型统计切片，而不是全局数组，修改后的多进程代码如下：

import numpy as np
import multiprocessing as mp
import itertools

# 假设count_types函数已正确实现，返回(n_bigram_types_chunk, row_types_chunk, col_types_chunk)
def count_types(chunk):
    # 原实现逻辑保持不变
    pass

def kn_probs(bigram_counts, type_counts, d=0.75):
    # 原函数逻辑保持不变
    n_bigram_types, row_types_counts, col_types_counts = type_counts

    kn_probs = np.zeros((len(bigram_counts), len(bigram_counts[0])), dtype=np.float16)

    for row_i, row in enumerate(bigram_counts):
        pre_word_count = np.sum(row)
        n_pre_word_followers = row_types_counts[row_i]
        kn_lambda = (d / pre_word_count) * n_pre_word_followers
        
        for col_i, bigram_count in enumerate(row):
            n_word_precedents = col_types_counts[col_i]
            kn_prob = (max(bigram_count-d, 0) / pre_word_count)\
                        + (kn_lambda * (n_word_precedents / n_bigram_types))
            kn_probs[row_i][col_i] = kn_prob

    return kn_probs

# 修改后的多进程处理逻辑
n_processes = 16
bigram_counts_chunks = np.array_split(bigram_counts, n_processes)
length = len(bigram_counts)
del(bigram_counts)

with mp.Pool(n_processes) as p:
    # 计算各子矩阵的类型统计
    chunks_type_counts = p.map(count_types, bigram_counts_chunks)
    
    # 合并全局统计项
    n_bigram_types = sum([tc[0] for tc in chunks_type_counts])
    # 保留每个子矩阵对应的行类型统计，不做全局拼接
    row_types_chunks = [tc[1] for tc in chunks_type_counts]
    # 合并全局列类型统计，改用numpy数组初始化提升性能
    col_types_counts = np.zeros(length, dtype=np.int64)
    for tc in chunks_type_counts:
        col_types_counts = np.add(col_types_counts, tc[2])
    
    # 为每个子矩阵准备专属的type_counts参数
    chunk_args = []
    for chunk, row_types in zip(bigram_counts_chunks, row_types_chunks):
        # 每个子矩阵的type_counts：全局n_bigram_types + 局部row_types + 全局col_types_counts
        chunk_args.append( (chunk, (n_bigram_types, row_types, col_types_counts)) )
    
    # 并行计算各子矩阵的概率
    bigram_probs_chunks = p.starmap(kn_probs, chunk_args)
    del(bigram_counts_chunks)

# 合并最终概率矩阵
bigram_probs = np.concatenate(bigram_probs_chunks)

关键修改点 #

1. 不再拼接全局row_types_counts，而是保留每个子矩阵对应的row_types_chunk
2. 为每个子矩阵单独构造type_counts参数，确保row_types_counts是当前子矩阵对应的局部切片，和函数内的局部行索引匹配
3. 把col_types_counts的初始化从列表改为np.zeros，减少类型转换的性能损耗

额外优化建议 #

• 考虑用np.float32代替np.float16：float16精度有限，可能在计算中出现数值溢出或精度丢失，导致异常值
• 如果矩阵规模极大，可考虑用共享内存传递大数组，减少多进程间的数据拷贝开销

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