PyTorch为了在不同进程间高效共享张量数据，提供了多种共享策略。当调用torch.multiprocessing.set_sharing_strategy("file_system")时，PyTorch会利用操作系统的文件系统作为后端来实现共享内存。这意味着，当张量在进程间传递时，它们的数据可能不会被复制，而是通过在/dev/shm（或类似位置，如/dev）创建文件来映射到共享内存区域。这些文件通常以torch_shm_为前缀命名，用于存储张量的实际数据。
这种机制的优点是避免了大数据量的复制，提高了效率。然而，其缺点在于，如果共享内存段没有被正确或及时地清理，这些torch_shm_文件就会持续存在，占用/dev/shm空间，最终可能导致：
- 共享内存耗尽： 系统无法再分配新的共享内存段。
- 文件描述符耗尽： 大量文件打开可能超出操作系统限制。
- 程序崩溃： 由于资源不足，PyTorch或底层系统调用失败。
用户提供的代码片段展示了典型的PyTorch多进程设置：
import torch
# 强制使用文件系统共享策略
torch.multiprocessing.set_sharing_strategy("file_system")
from glob import glob
import os # Added for path manipulation
def processData(fname):
# 模拟张量处理，返回一个张量列表
# 实际应用中，这里会进行复杂的张量操作
return [torch.randn(10, 10) for _ in range(5)] # 示例：返回5个张量
if __name__ == "__main__":
datadir = "./temp_data" # 使用相对路径方便测试
# 创建一些示例文件
os.makedirs(datadir, exist_ok=True)
for i in range(25000): # 模拟大量数据点
with open(os.path.join(datadir, f"data_{i}.txt"), "w") as f:
f.write(f"This is data point {i}")
processedData = []
# 设置多进程启动方法为'spawn'，确保进程隔离
torch.multiprocessing.set_start_method("spawn", force=True) # force=True避免重复设置错误
print("Starting multiprocessing pool...")
p = torch.multiprocessing.Pool(8)
for i, result in enumerate(p.imap_unordered(processData, glob(datadir+"/*.txt"))):
processedData.extend(result)
if i % 1000 == 0:
print(f"Processed {i+1} files. Current /dev/shm usage (approx): {len(os.listdir('/dev/shm'))} files")
p.close()
p.join()
print("Multiprocessing finished.")
# 清理临时文件
for f in glob(datadir+"/*.txt"):
os.remove(f)
os.rmdir(datadir) 在上述代码中，processData函数返回一个张量列表，这些张量通过共享内存机制传递回主进程。当数据量达到一定程度时（如示例中的20,000个数据点），/dev/shm中的torch_shm_文件会迅速累积，最终导致资源耗尽。
2. 问题根源：共享内存段的生命周期管理尽管Pool对象在调用p.close()和p.join()后会尝试清理其管理的进程和资源，但对于file_system策略下的共享内存文件，其生命周期管理可能不如预期。以下是可能导致文件堆积的原因：
PIA
全面的AI聚合平台，一站式访问所有顶级AI模型
226
查看详情
- 引用计数问题： PyTorch内部可能使用引用计数来管理共享内存段。如果主进程或子进程对共享内存中的张量仍持有引用，即使进程已结束，对应的共享内存文件也可能不会立即被释放。
- 异常终止： 如果子进程在处理过程中崩溃或被强制终止，它们可能没有机会正常释放所持有的共享内存资源。
- 操作系统延迟： 即使所有引用都已释放，操作系统也可能不会立即删除这些文件，尤其是在高负载情况下。
- 设计权衡： 为了性能，PyTorch可能选择在某些情况下延迟清理，或者依赖于系统重启来清除这些临时文件。
用户尝试了分块处理和为每个分块创建新的Pool对象，以及del返回结果，但问题依然存在。这表明问题可能不在于Python对象的垃圾回收，而在于底层共享内存文件的生命周期管理。
3. 优化策略与注意事项鉴于PyTorch在file_system策略下共享内存清理的潜在问题，以下是一些缓解策略和专业建议：
3.1 优化数据传输与处理- 减少共享数据量： 重新评估processData函数返回的数据量。如果返回的是大量中间张量，考虑是否可以在子进程中完成更多处理，只返回最终的、精简的结果。
- 非共享数据传输： 如果可能，避免通过共享内存传输张量。例如，子进程可以将处理结果序列化到磁盘，然后主进程再从磁盘加载。这会增加I/O开销，但可以避免共享内存问题。
- 使用Queue或Pipe传递少量元数据： 对于少量非张量数据，可以使用torch.multiprocessing.Queue或Pipe进行进程间通信，这通常不会导致共享内存文件问题。
- 显式释放张量： 尽管用户已尝试del，但再次强调，在不再需要张量时，显式将其设置为None并调用gc.collect()有助于Python的垃圾回收。虽然这不直接清理/dev/shm文件，但可以帮助释放Python层面的引用。
- 限制并发数量： 适当减少Pool的并发进程数（例如，从8减少到4），可以降低/dev/shm的瞬时压力。
- 分批处理与周期性重启Pool： 用户尝试的这种方法是有效的思路。如果每次处理大量数据后重启Pool仍然无法清理，可能需要结合系统级清理。确保在每次Pool结束后，所有子进程都已完全终止。
在绝大多数情况下，不建议手动清理/dev/shm中的torch_shm_文件，尤其是当PyTorch进程仍在运行时。这样做可能导致数据损坏或程序崩溃。然而，在开发和调试阶段，或者在程序完全停止后，作为最后的手段，可以考虑以下方法：
- 系统重启： 最简单也是最安全的系统
以上就是PyTorch多进程共享内存管理：解决/dev/shm文件堆积问题的详细内容，更多请关注知识资源分享宝库其它相关文章！
相关标签： python 操作系统 大数据 app 后端 ai red Python 堆 并发 对象 pytorch 大家都在看： python怎么检查一个键是否存在于字典中_python字典键存在性检查 Python怎么实现一个上下文管理器_Python上下文管理器协议实现 python中怎么给函数设置默认参数_Python函数默认参数设置方法 python中怎么测量一段代码的执行时间？ python怎么创建一个虚拟环境_python虚拟环境创建与使用教程
发表评论:
◎欢迎参与讨论，请在这里发表您的看法、交流您的观点。