tqdmprocess_map大数组优化技巧

从现在开始，我们要努力学习啦！今天我给大家带来《tqdm process_map传递大数组：共享内存优化方法》，感兴趣的朋友请继续看下去吧！下文中的内容我们主要会涉及到等等知识点，如果在阅读本文过程中有遇到不清楚的地方，欢迎留言呀！我们一起讨论，一起学习！

在使用`tqdm.contrib.concurrent.process_map`进行并行处理时，直接将大型数组作为函数参数传递可能因数据复制导致`MemoryError`。本教程将介绍如何利用`multiprocessing.Array`创建共享内存，使多个进程能够高效访问同一份大型数组数据，避免昂贵的内存复制，从而优化内存使用并实现健壮的并行计算。

1. 并行处理大型数组的挑战

tqdm.contrib.concurrent.process_map提供了一种便捷的方式来并行化函数调用，并带有进度条显示，这对于处理耗时任务非常有用。然而，当并行函数需要访问一个大型数组（例如一个NumPy数组）时，直接传递这个数组作为函数参数可能会遇到内存效率问题。

考虑以下场景：有一个大型数组B，以及一个需要对B的特定部分进行操作的函数test(a, B)。如果尝试通过构造一个包含B的元组列表（例如agrid = [(0, B), (1, B), ...]）来传递参数给process_map，multiprocessing模块在将这些参数发送给子进程时，会尝试序列化并复制B。对于非常大的B，这种复制会导致显著的内存开销，甚至可能引发MemoryError，因为每个子进程都会获得B的一个独立副本。

问题的核心在于，我们希望所有子进程都能访问同一个大型数组B，而不是各自拥有一个副本。

2. 解决方案：利用multiprocessing.Array实现共享内存

为了解决大型数组的内存复制问题，Python的multiprocessing模块提供了Array类，它允许在进程间共享原始数据类型的数组。通过将大型数组存储在共享内存中，子进程可以直接访问这块内存，而无需进行昂贵的数据复制。

其基本原理如下：

1. 创建共享内存数组： 使用multiprocessing.Array分配一块指定类型和大小的内存区域。ctypes模块用于定义数组中元素的类型（例如ctypes.c_double表示双精度浮点数）。
2. 转换为NumPy视图： multiprocessing.Array返回的是一个低级的C类型数组对象。为了方便地使用NumPy的强大功能，我们可以通过np.frombuffer(shared_array.get_obj()).reshape(...)将其转换为一个NumPy数组的视图。这个视图直接操作共享内存，而不是复制数据。
3. 进程访问： 在主进程中初始化共享数组并填充数据后，子进程可以通过访问主进程中设置的（通常是全局的）multiprocessing.Array实例，然后将其转换为NumPy视图来访问共享数据。

3. 实践示例：共享大型NumPy数组

下面是一个详细的示例，展示了如何使用multiprocessing.Array与tqdm.contrib.concurrent.process_map结合，高效地处理大型NumPy数组。

import ctypes
from multiprocessing import Array
from time import sleep

import numpy as np
from tqdm.contrib.concurrent import process_map

# 定义一个全局变量来持有共享数组的引用
# 注意：在多进程环境中，全局变量的赋值需要在if __name__ == "__main__": 块内
# 并且子进程会继承父进程的全局变量副本，但对于multiprocessing.Array，
# 它们会指向同一块共享内存区域。
B_shared = None 

# 数组的维度
N = 1_000

def test(a):
    """
    工作函数：操作共享内存中的大型数组。
    这个函数只接收一个索引 'a' 作为参数。
    """
    # 将共享内存对象转换为NumPy数组视图
    # B_shared 必须在进程启动时被正确初始化
    arr = np.frombuffer(B_shared.get_obj()).reshape((N, N))

    # 模拟复杂的计算
    sleep(0.1) # 减少睡眠时间以加快示例运行

    # 如果需要写入共享数组，必须使用锁来同步访问
    # 例如：
    # with B_shared.get_lock():
    #     arr[a, a] = some_new_value

    # 返回数组中特定位置的值
    return arr[a, a]


if __name__ == "__main__":
    # 1. 初始化共享内存数组
    # ctypes.c_double 指定数组元素类型为双精度浮点数
    # N * N 是数组的总元素数量
    B_shared = Array(ctypes.c_double, N * N)

    # 2. 将共享内存转换为NumPy数组视图，并填充数据
    # arr 是一个NumPy数组，但它的数据存储在B_shared管理的共享内存中
    arr_view = np.frombuffer(B_shared.get_obj()).reshape((N, N))
    arr_view[:] = np.random.uniform(size=(N, N)) # 填充随机数据

    print(f"原始数组B_shared的前几个元素:\n{arr_view[:2, :2]}")

    # 3. 定义要迭代的参数列表
    agrid = [0, 1, 2, 3] # 假设我们想对这些索引进行操作

    # 4. 使用 process_map 进行并行计算
    # test 函数将会在每个子进程中执行
    # max_workers 控制并行进程数
    # chunksize 控制每次发送给子进程的任务块大小
    parallel_results_tqdm = process_map(
        test,
        agrid,
        max_workers=2,
        chunksize=1,
    )

    print(f"\n并行计算结果: {parallel_results_tqdm}")
    print(f"计算完成后B_shared的前几个元素 (如果未写入，则与原始相同):\n{arr_view[:2, :2]}")

代码解析与注意事项：

• B_shared = None： 在全局作用域声明B_shared，使其在if __name__ == "__main__":块中被初始化后，其引用能被子进程继承。子进程启动时，会复制父进程的全局变量环境，但对于multiprocessing.Array对象，这个引用会指向同一块物理共享内存。
• if __name__ == "__main__":： 这是Python多进程编程的关键。所有创建子进程的代码（包括process_map）都必须放在这个保护块内，以避免在Windows系统上出现递归创建进程的问题，并确保代码的正确执行。
• *`Array(ctypes.c_double, N N)：** 这行代码在主进程中分配了一块足够存储N * N`个双精度浮点数的共享内存。
• np.frombuffer(B_shared.get_obj()).reshape((N, N))： 这是将共享内存转换为NumPy数组视图的关键步骤。B_shared.get_obj()返回底层的ctypes数组对象，np.frombuffer则基于这个缓冲区创建NumPy数组，reshape将其塑形为所需的二维数组。
• test(a)函数： 现在test函数只接收一个索引a作为参数。在函数内部，它通过全局变量B_shared访问共享内存，并将其转换为NumPy视图进行操作。
• 写入共享内存： 如果子进程需要修改共享数组中的数据，必须使用B_shared.get_lock()获取一个锁。这可以防止多个进程同时写入同一块内存区域，从而避免数据竞争和不一致性。在示例代码中，写入操作被注释掉了，但如果需要，请务必使用锁。

4. 总结

通过利用multiprocessing.Array创建共享内存，我们成功解决了在tqdm.contrib.concurrent.process_map中传递大型数组参数时可能遇到的MemoryError问题。这种方法避免了昂贵的数据复制，显著提高了内存效率，并使得并行处理大型数据集成为可能。

关键要点包括：

• 使用multiprocessing.Array分配共享内存。
• 将共享内存转换为NumPy数组视图以便于操作。
• 将共享数组的引用作为全局变量（或通过initializer和initargs）传递给子进程。
• 始终将多进程代码置于if __name__ == "__main__":保护块内。
• 如果子进程需要修改共享数组，务必使用multiprocessing.Lock进行同步，以防止数据竞争。

掌握这种技术，将能更高效地利用多核CPU资源，处理计算密集型且涉及大型数据集的Python任务。

今天带大家了解了的相关知识，希望对你有所帮助；关于文章的技术知识我们会一点点深入介绍，欢迎大家关注golang学习网公众号，一起学习编程~