如何并行化Python嵌套循环以加速Wasserstein距离计算？

问题 #

使用Scipy计算wasserstein_distance时，因data[row][column]单条数据量达6万-30万条，四层嵌套循环的串行计算耗时极长。期望将计算任务按行分配给多核处理器（如第一个处理器处理(0,0)至(0,23)所有像素与其他像素的距离，第二个处理器处理(1,0)至(1,23)的像素），但尝试基于ThreadPoolExecutor的并行实现后，遍历futures的阶段仍耗时很久，寻求可行的并行优化方案。

原串行代码 #

def compute_wasserstein_distances(tflite, data, rows, columns, output_filename):

    # Initialize an empty dictionary to store wasserstein distances
    dist_similarity = {}
    
    # Iterate through rows and columns of the grid
    for row in tqdm(range(0, rows)):              
        for column in tqdm(range(0, columns), leave= False):
            for r in range(0, rows):
                for c in range(0, columns):
                    # Calculate the Wasserstein distance between the distributions
                    distance = wasserstein_distance(data[row][column], data[r][c])
                    
                    # Store the distance in the dictionary with corresponding indices
                    dist_similarity[((row, column), (r, c))] = distance

    print(f"[INFO] Saving pixel distribution to file ...")
    
    # Save the dist_similarity dictionary to a pickle file
    with open(f'{output_filename}/{tflite}-DistSimilarity.pkl', 'wb') as file:
        pickle.dump(dist_similarity, file)
    
    return dist_similarity 

尝试的并行代码 #

import pickle
import concurrent.futures
from tqdm import tqdm
from scipy.stats import wasserstein_distance

def calculate_distance(row, column, r, c, data):
    return ((row, column), (r, c), wasserstein_distance(data[row][column], data[r][c]))

def compute_wasserstein_distances_parallel(tflite, data, rows, columns, output_filename):
    # Initialize an empty dictionary to store Wasserstein distances
    dist_similarity = {}

    # Create a ThreadPoolExecutor
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = []

        # Iterate through rows and columns of the grid
        for row in tqdm(range(rows)):
            for column in tqdm(range(columns), leave=False):
                for r in range(rows):
                    for c in range(columns):
                        # Submit the task to the executor
                        futures.append(executor.submit(calculate_distance, row, column, r, c, data))

        # Wait for all tasks to complete
        concurrent.futures.wait(futures)

        # Retrieve results from futures
        for future in tqdm(futures):
            indices, distance = future.result()[:2], future.result()[2]
            dist_similarity[indices] = distance

    print(f"[INFO] Saving pixel distribution to file ...")

    # Save the dist_similarity dictionary to a pickle file
    with open(f'{output_filename}/{tflite}-DistSimilarity.pkl', 'wb') as file:
        pickle.dump(dist_similarity, file)

    return dist_similarity

优化方案 #

1. 替换线程池为进程池 #

wasserstein_distance属于CPU密集型任务，Python的GIL会限制线程池的多核利用率，改用ProcessPoolExecutor可真正利用多处理器资源。

2. 优化任务粒度 #

避免为每个(r,c)单独提交任务，改为以(row, column)为单位提交任务——每个任务负责计算当前像素与所有其他像素的距离，大幅减少任务提交和结果处理的开销。

3. 减少重复计算 #

Wasserstein距离满足对称性：d(a,b) = d(b,a)，只需计算上三角（或下三角）区域的距离，再对称填充结果，可减少约50%的计算量。

4. 优化结果存储 #

用numpy数组替代字典存储距离矩阵，数组的读写和内存效率远高于字典，最后可按需转换为字典或直接保存数组。

优化后的代码示例 #

import pickle
import numpy as np
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm
from scipy.stats import wasserstein_distance

def process_pixel(args):
    row, column, data, rows, columns = args
    distances = np.zeros((rows, columns))
    for r in range(rows):
        for c in range(columns):
            distances[r, c] = wasserstein_distance(data[row][column], data[r][c])
    return row, column, distances

def compute_wasserstein_distances_optimized(tflite, data, rows, columns, output_filename):
    # 初始化距离矩阵
    dist_matrix = np.zeros((rows, columns, rows, columns))
    
    # 准备任务参数
    tasks = [(row, col, data, rows, columns) for row in range(rows) for col in range(columns)]
    
    with ProcessPoolExecutor() as executor:
        # 提交任务并处理结果
        for row, col, distances in tqdm(executor.map(process_pixel, tasks), total=len(tasks)):
            dist_matrix[row, col] = distances
    
    # 转换为字典（若需保持原格式）
    dist_similarity = {}
    for row in range(rows):
        for col in range(columns):
            for r in range(rows):
                for c in range(columns):
                    dist_similarity[((row, col), (r, c))] = dist_matrix[row, col, r, c]
    
    print(f"[INFO] Saving pixel distribution to file ...")
    # 可选择直接保存numpy数组（更高效）
    # np.save(f'{output_filename}/{tflite}-DistSimilarity.npy', dist_matrix)
    with open(f'{output_filename}/{tflite}-DistSimilarity.pkl', 'wb') as file:
        pickle.dump(dist_similarity, file)
    
    return dist_similarity

额外优化建议 #

• 若data是列表格式，先转换为numpy数组，减少内存拷贝开销，且scipy对numpy数组的处理效率更高。
• 若内存允许，可将data作为全局变量在子进程中初始化，避免进程间大量数据传递。
• 利用对称性优化：只计算row*columns + col <= r*columns + c的部分，再填充对称值，进一步减少计算量。

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