PyTorch DataLoader含自定义Transform时num_workers>0程序挂起问题

核心原因 #

当num_workers>0时，DataLoader会启动子进程加载数据，自定义Transform如果未实现多进程兼容，会导致进程间通信阻塞，进而无限挂起。哪怕你的MedianFilter是恒等函数，只要类的实现不符合PyTorch多进程下的序列化要求，就会触发问题。

常见的不兼容点：

• Transform类的__init__或__call__中使用了无法被pickle序列化的对象（比如非标准库资源、未正确实现序列化的类实例）
• 类定义依赖父进程的全局状态或动态资源，子进程无法正确复制这些状态

解决方法 #

针对你的场景，有几种可靠的修复方式：

1. 确保Transform类可被正确序列化 #

PyTorch多进程加载依赖pickle序列化Transform对象并传递给子进程，需保证自定义类能被正确处理：

import torchvision.transforms as transforms

class MedianFilter(transforms.Transform):
    def __init__(self):
        super().__init__()

    def __call__(self, img):
        # 恒等操作直接返回输入
        return img

若类中有其他属性，需确保这些属性均为可序列化类型（避免使用lambda、未序列化的句柄等）。

2. 用函数式Transform替代类实现 #

函数式实现天生更易被pickle序列化，可规避类带来的潜在问题：

def median_filter(img):
    return img

# 在transform组合中使用
transform = transforms.Compose([
    median_filter,
    transforms.ToTensor()
])

3. 调整DataLoader的多进程启动方式 #

Windows系统默认用spawn启动子进程，Linux/macOS默认用fork。可尝试显式指定更稳定的启动方式：

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    multiprocessing_context='forkserver'  # 或'spawn'
)

forkserver会创建独立服务器进程生成子进程，能避免父进程状态复制的问题。

4. 避免在Transform中依赖全局变量或动态资源 #

若MedianFilter需依赖资源，将初始化逻辑放在__call__方法中，确保每个子进程独立初始化：

class MedianFilter(transforms.Transform):
    def __init__(self):
        super().__init__()
        self.resource = None

    def __call__(self, img):
        # 每个子进程独立初始化资源
        if self.resource is None:
            # 资源初始化逻辑
            pass
        return img

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