PyTorch多进程共享内存溢出：深度解析与管理策略

理解PyTorch共享内存机制

pytorch为了高效地在不同进程间共享张量数据，利用了操作系统提供的共享内存机制。当一个张量被创建并设置为可共享时，其底层数据会被存储在共享内存区域。在linux系统中，这通常表现为在/dev/shm（共享内存文件系统）下创建以torch_shm_为前缀的临时文件或目录。

当使用torch.multiprocessing模块（特别是Pool）进行并行计算时，如果子进程返回的张量或包含张量的对象需要被主进程接收，PyTorch会尝试通过共享内存来避免昂贵的数据拷贝。torch.multiprocessing.set_sharing_strategy("file_system")策略指示PyTorch使用文件系统作为共享内存的后端，这对于处理大量张量或避免文件描述符限制非常有用，但同时也增加了对/dev/shm空间的需求。

问题情景分析与代码示例

在并行处理大量数据时，开发者可能会遇到PyTorch多进程导致共享内存溢出的问题。具体表现为，当处理的数据点达到一定数量后，程序因无法分配更多共享内存而崩溃。通过检查/dev/shm目录，会发现大量未被清理的torch_shm_文件持续堆积。即使尝试了分块处理、为每个分块创建新的Pool对象以及显式删除返回结果，问题可能依然存在。

以下是一个典型的代码示例，展示了可能导致共享内存溢出的场景：

import torch
import torch.multiprocessing as mp
from glob import glob
import os
import time

# 由于操作系统文件描述符限制，通常需要使用 "file_system" 策略
mp.set_sharing_strategy("file_system")

def processData(fname):
    """
    模拟数据处理函数，从文件中读取数据并生成一个张量列表。
    这里的张量会被共享给主进程。
    """
    # 实际场景中，这里会进行复杂的张量操作
    # 示例：创建一个包含5个随机张量的列表
    return [torch.randn(100, 100, dtype=torch.float32) for _ in range(5)]

if __name__ == "__main__":
    # 设置启动方法为 'spawn'，确保跨平台兼容性和安全性
    # force=True 可以在脚本多次运行时避免重复设置的错误
    mp.set_start_method("spawn", force=True)

    # 替换为实际的数据目录，例如创建一个虚拟目录和文件
    # for i in range(20000):
    #     with open(f"/tmp/data/file_{i}.txt", "w") as f:
    #         f.write(f"data for file {i}")
    datadir = "/path/to/your/data/directory" # 请替换为实际路径

    processedData = []

    print(f"程序启动时 /dev/shm 中的 torch_shm_ 文件数量: {len([f for f in os.listdir('/dev/shm') if f.startswith('torch_shm_')])}")

    # 创建一个进程池
    with mp.Pool(processes=8) as pool: # 使用 with 语句确保进程池被正确关闭
        file_paths = glob(datadir + "/*.txt")
        print(f"发现 {len(file_paths)} 个文件待处理。")

        # 使用 imap_unordered 异步获取结果
        for i, result_list in enumerate(pool.imap_unordered(processData, file_paths)):
            processedData.extend(result_list)

            # 定期检查 /dev/shm 状态
            if (i + 1) % 100 == 0:
                shm_files_count = len([f for f in os.listdir('/dev/shm') if f.startswith('torch_shm_')])
                print(f"已处理 {i+1} 个文件，当前 /dev/shm 中的 torch_shm_ 文件数量: {shm_files_count}")
                # 可以在这里添加检查，如果文件数量异常，则提前退出或记录警告

    # 进程池在 with 语句结束后会自动 close 和 join

    shm_files_count_final = len([f for f in os.listdir('/dev/shm') if f.startswith('torch_shm_')])
    print(f"所有文件处理完毕后 /dev/shm 中的 torch_shm_ 文件数量: {shm_files_count_final}")

    # 此时 processedData 列表中包含了所有子进程返回的张量
    # 如果这些张量不被及时清理，其对应的共享内存区域将持续占用
    # 显式解除引用，有助于垃圾回收
    del processedData
    # 强制进行Python垃圾回收，但这不直接作用于共享内存文件
    import gc
    gc.collect()

    # 再次检查，看是否有变化 (通常不会立即清理共享内存文件)
    shm_files_count_after_del = len([f for f in os.listdir('/dev/shm') if f.startswith('torch_shm_')])
    print(f"显式删除引用后 /dev/shm 中的 torch_shm_ 文件数量: {shm_files_count_after_del}")

尽管代码中使用了mp.Pool的with语句（或显式的pool.close()和pool.join())来确保进程池被正确关闭，但共享内存文件可能并未及时清理。这通常是因为PyTorch的共享内存管理器在某些情况下未能及时回收资源，或者主进程仍然持有对这些共享内存区域的隐式引用。

共享内存溢出排查与诊断

要诊断共享内存溢出问题，最直接且有效的方法是监控/dev/shm目录。在Linux终端中，可以使用以下命令：

ls -l /dev/shm | grep torch_shm_ | wc -l

该命令会统计当前/dev/shm目录下以torch_shm_开头的文件的数量。在程序运行期间定期执行此命令，可以观察文件数量是否持续增长。如果文件数量快速增加且不下降，则强烈表明存在共享内存泄漏。

此外，还可以使用df -h /dev/shm命令来查看/dev/shm文件系统的实际使用情况，判断是否接近满载。

ToonMe

一款风靡Instagram的软件，一键生成卡通头像

下载

现有解决方案与注意事项

当遇到PyTorch多进程共享内存溢出时，可以考虑以下解决方案，并注意其局限性和潜在风险。

1. 系统重启

最直接有效的临时解决方案是重启操作系统或计算环境。系统重启会清空/dev/shm目录，从而释放所有被占用的共享内存资源。

• 优点： 简单、彻底，能够快速解决当前的内存溢出问题。
• 局限性： 这种方法在开发和测试阶段可行，但在生产环境或需要长时间运行、高可用性的任务中显然不切实际。它治标不治本，不能解决根本的内存管理问题。

2. 手动清理的风险

理论上，可以通过手动删除/dev/shm目录下的torch_shm_文件来释放空间。例如，使用以下命令：

sudo rm -rf /dev/shm/torch_shm_*

• 严重警告： 直接在/dev或/dev/shm目录下使用rm -rf命令是极其危险的操作。如果命令中不慎出现空格（例如rm -rf torch *），可能导致删除/dev目录下所有文件，甚至整个文件系统，造成不可逆的数据丢失和系统损坏。此外，如果程序正在运行并依赖这些共享内存文件，强行删除可能导致程序崩溃或数据损坏。因此，强烈不建议在生产环境或不熟悉系统操作的情况下执行此类命令。即使要手动清理，也应精确指定路径和前缀，并仔细检查命令，且仅在确认没有PyTorch进程正在使用这些文件时进行。

PyTorch共享内存管理策略探讨

为了从根本上解决或缓解PyTorch多进程共享内存溢出问题，需要深入理解并优化共享内存的使用方式。

1. 优化数据传输

• 减少共享张量数量： 仔细检查processData函数返回的数据。如果返回的是大量小型张量列表，每个张量都可能在共享内存中创建独立的区域。考虑是否可以：

  • 合并张量： 将多个小张量合并成一个大张量返回，减少共享内存对象的数量。
  • 内部处理，只返回结果： 在子进程内部完成大部分处理，只返回最终的、聚合的结果（例如，标量、小列表、或经过降维的张量），而非原始张量或中间计算结果。
  • 序列化传输： 如果数据量巨大且不适合共享内存，可以考虑将子进程处理后的数据序列化（例如，使用pickle、numpy.save、torch.save），然后通过队列或文件传输序列化数据的路径，在主进程中反序列化加载。这会增加CPU开销，但能避免共享内存问题。

• 显式管理张量生命周期： 尽管PyTorch的共享内存管理器会尝试清理，但如果主进程对子进程返回的张量（或包含它们的列表）保持引用，这些共享内存区域就不会被释放。确保在不再需要时，显式地解除对这些张量的引用，让Python的垃圾回收机制有机会清理它们。例如，可以定期将processedData列表中的数据持久化到磁盘，然后清空列表，或者在处理完一批数据后，使用del关键字显式删除不再需要的张量引用。

  # 示例：定期清理 processedData 列表
  batch_size = 1000
  for i, result_list in enumerate(pool.imap_unordered(processData, file_paths)):
      processedData.extend(result_list)
      if (i + 1) % batch_size == 0:
          # 在这里可以将 processedData 写入磁盘或进行其他聚合操作
          # 然后清空列表，解除对张量的引用
          del processedData[:] # 清空列表内容
          print(f"已处理 {i+1} 个文件，并清空 processedData 列表。")

2. 进程生命周期管理

• 确保进程池正确关闭： 使用with mp.Pool(...) as pool:语句是管理进程池生命周期的最佳实践。它确保了在代码块结束时，无论是否发生异常，进程池都会被正确关闭（close()）并等待所有子进程完成（join()），从而允许PyTorch有机会清理它们持有的共享内存资源。
• **分