如何用VizTracer追踪PyTorch DataLoader的Worker进程？

核心问题在于VizTracer默认仅识别标准库的multiprocessing.Process，而PyTorch DataLoader使用的是torch.multiprocessing.Process子类，导致CLI模式下无法自动追踪worker进程。以下是两种无需大量修改代码、直接通过CLI实现完整追踪的可行方案：

方案一：预加载补丁脚本（零修改主代码） #

创建一个补丁脚本patch_torch_mp.py，让VizTracer识别PyTorch的进程类并自动追踪worker：

import torch.multiprocessing as mp
from viztracer import get_tracer

def patch_viztracer_torch_mp():
    tracer = get_tracer()
    if tracer is not None:
        # 将PyTorch的Process类加入VizTracer的可追踪进程列表
        tracer.trace_multiprocess_types.add(mp.Process)
        
        # 给PyTorch进程添加启动钩子，确保子进程自动初始化追踪
        original_start = mp.Process.start
        def patched_start(self):
            self._viztracer_trace = True
            original_start(self)
        mp.Process.start = patched_start

patch_viztracer_torch_mp()

然后通过VizTracer的--preload参数加载补丁后运行主脚本：

python -m viztracer --trace_multiprocess --preload patch_torch_mp.py myscript.py

该方案完全不需要修改主代码，所有追踪逻辑通过补丁注入，且VizTracer会自动合并所有进程的追踪数据到单个报告文件。

方案二：主代码添加极简兼容逻辑（仅2-3行改动） #

如果可以接受主代码的极小改动，在脚本开头加入以下条件执行的代码（仅当VizTracer运行时生效，不影响正常逻辑）：

import sys
if "viztracer" in sys.modules:
    import torch.multiprocessing as mp
    from viztracer import get_tracer
    tracer = get_tracer()
    if tracer:
        tracer.trace_multiprocess_types.add(mp.Process)

之后直接用标准CLI命令运行即可：

python -m viztracer --trace_multiprocess myscript.py

关键说明 #

• 避免直接替换torch.multiprocessing.Process为标准库进程类：这会破坏PyTorch针对CUDA张量共享、进程初始化的特殊逻辑，导致数据加载异常。
• --trace_multiprocess参数的作用：自动在子进程启动VizTracer追踪，并将所有进程的追踪数据合并到同一个输出文件，无需手动合并。
• 适配PyTorch的spawn启动模式：上述方案均处理了spawn模式下的追踪初始化问题，覆盖DataLoader的默认进程启动方式。

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