PyTorch自定义AudioDataset多进程加载报错与性能优化问题咨询

报错问题修复 #

该报错是PyTorch多进程DataLoader使用spawn启动模式（Windows/macOS默认，Linux部分环境也会启用）时的序列化问题：子进程启动时需要从__main__模块导入AudioDataset类，当前类的定义位置或主逻辑写法不满足序列化要求，可按以下步骤修复：

• 把AudioDataset类的定义放到代码顶级作用域，不要嵌套在其他函数、类内部。
• 所有主执行逻辑（数据集实例化、DataLoader实例化、训练逻辑等）都用if __name__ == "__main__"包裹，示例：

if __name__ == "__main__":
    # 路径、目标列表初始化代码
    dataset = AudioDataset(paths_list, targets)
    dataloader = DataLoader(dataset, batch_size=8, num_workers=10)
    # 后续训练逻辑

• 如果你在Jupyter Notebook等交互式环境运行，spawn模式不支持顶级类序列化，要么切换为Python脚本运行，要么在Linux环境下设置多进程启动方式为fork，在代码开头添加：

import multiprocessing
multiprocessing.set_start_method('fork', force=True)

性能优化方案 #

当前初始化耗时久的核心原因是你把所有音频预处理逻辑放在__init__中串行执行，即使开启DataLoader的num_workers也不会生效，因为预处理在主进程初始化阶段就已经全部跑完了，num_workers仅负责__getitem__阶段的加载逻辑。可选择以下两种方案优化：

方案1：懒加载+多Worker预处理（改动最小，推荐） #

把预处理逻辑从__init__移到__getitem__中，让DataLoader的多Worker并行处理每个样本的音频，初始化可以瞬间完成，整体处理速度提升和Worker数成正比：
修改后的Dataset代码：

class AudioDataset(Dataset):
    def __init__(self, paths_list, targets, preprocess=preprocess_fn):
        self.preprocess = preprocess
        self.paths_list = paths_list
        self.targets = targets
        self.cache = {} # 内存缓存，内存充足时开启，避免每个epoch重复处理

    def __len__(self):
        return len(self.paths_list)

    def __getitem__(self, idx):
        if idx in self.cache:
            return self.cache[idx]
        p = self.paths_list[idx]
        target = self.targets[idx]
        audio_1 = self.preprocess(p[0])
        audio_2 = self.preprocess(p[1])
        res = ([audio_1, audio_2], target)
        self.cache[idx] = res
        return res

方案2：初始化阶段并行预处理+本地缓存（适合数据集固定、反复使用的场景） #

如果数据集不会频繁变更，可以在初始化阶段用多进程并行处理所有音频，还可以将处理好的特征存到本地，下次运行直接读缓存，不需要重复处理：

from concurrent.futures import ProcessPoolExecutor
import torch
import os

class AudioDataset(Dataset):
    def __init__(self, paths_list, targets, preprocess=preprocess_fn, cache_path=None):
        self.preprocess = preprocess
        self.features = []
        self.labels = []
        self.paths_list = paths_list
        self.targets = targets
        
        # 有本地缓存直接读缓存，跳过预处理
        if cache_path and os.path.exists(cache_path):
            cache = torch.load(cache_path)
            self.features = cache["features"]
            self.labels = cache["labels"]
        else:
            self._init_dataset()
            # 存本地缓存
            if cache_path:
                torch.save({"features": self.features, "labels": self.labels}, cache_path)
    
    def _process_single(self, p, target):
        audio_1 = self.preprocess(p[0])
        audio_2 = self.preprocess(p[1])
        return [audio_1, audio_2], target

    def _init_dataset(self):
        with ProcessPoolExecutor(max_workers=10) as executor:
            futures = []
            for p, target in zip(self.paths_list, self.targets):
                futures.append(executor.submit(self._process_single, p, target))
            for future in tqdm(futures):
                feat, label = future.result()
                self.features.append(feat)
                self.labels.append(label)

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        sample = self.features[idx]
        label = self.labels[idx]
        return sample, label

内容的提问来源于stack exchange，提问作者Muhammad Anas Raza