Numpy memmap搭配PyTorch时num_workers>1致工作站冻结重启的解决办法咨询

1. 调整MemMap模式与Dataset初始化逻辑 #

你的代码中使用mode='r+'打开MemMap，但如果仅需读取数据，改为只读模式mode='r'可减少进程间的资源竞争与锁冲突。同时，避免在父进程提前创建MemMap对象——当num_workers>1时，每个子进程会复制父进程的MemMap实例，可能导致重复内存映射或资源过载。

修改后的Dataset实现，将文件路径传入，在子进程中按需打开MemMap：

from torch.utils.data import DataLoader, Dataset
import json
import numpy as np

class LargeFileDataset(Dataset):
    def __init__(self, x_file_path, y_file_path):
        self.x_path = x_file_path
        self.y_path = y_file_path
        
        # 预读取配置获取数据维度，避免重复IO
        with open(x_file_path + '.conf', 'r') as f:
            x_config = json.load(f)
        self.total_samples = x_config['shape'][0]
        self.x_shape = tuple(x_config['shape'])
        self.x_dtype = x_config['dtype']
        
        with open(y_file_path + '.conf', 'r') as f:
            y_config = json.load(f)
        self.y_shape = tuple(y_config['shape'])
        self.y_dtype = y_config['dtype']

    def __getitem__(self, idx):
        # 在每个样本读取时打开MemMap（或改为在worker_init_fn中全局初始化）
        x_mem = np.memmap(self.x_path, mode='r', shape=self.x_shape, dtype=self.x_dtype)
        y_mem = np.memmap(self.y_path, mode='r', shape=self.y_shape, dtype=self.y_dtype)
        return x_mem[idx], y_mem[idx]

    def __len__(self):
        return self.total_samples

2. 使用worker_init_fn统一初始化子进程MemMap #

为避免每个__getitem__重复打开文件，可通过worker_init_fn让每个子进程仅初始化一次MemMap，提升效率并减少资源占用：

# 全局字典存储子进程的MemMap对象
worker_shared_data = {}

def worker_init(worker_id):
    # 从全局字典获取文件路径
    x_path = worker_shared_data['x_path']
    y_path = worker_shared_data['y_path']
    
    with open(x_path + '.conf', 'r') as f:
        x_config = json.load(f)
    worker_shared_data['x'] = np.memmap(x_path, mode='r', shape=tuple(x_config['shape']), dtype=x_config['dtype'])
    
    with open(y_path + '.conf', 'r') as f:
        y_config = json.load(f)
    worker_shared_data['y'] = np.memmap(y_path, mode='r', shape=tuple(y_config['shape']), dtype=y_config['dtype'])

class WorkerInitDataset(Dataset):
    def __getitem__(self, idx):
        return worker_shared_data['x'][idx], worker_shared_data['y'][idx]

    def __len__(self):
        return worker_shared_data['x'].shape[0]

# 构建DataLoader时传入worker_init_fn
def my_dataloader(x_path, y_path, batch_size):
    worker_shared_data['x_path'] = x_path
    worker_shared_data['y_path'] = y_path
    dset = WorkerInitDataset()
    return DataLoader(
        dataset=dset,
        batch_size=batch_size,
        num_workers=4,  # 可根据硬件调整
        pin_memory=True,
        worker_init_fn=worker_init
    )

3. 切换为IterableDataset处理超大文件 #

对于远超内存的数据集，IterableDataset比常规Dataset更适合——它无需预存所有索引，可按需迭代读取，避免内存过载。同时支持多进程分片处理：

from torch.utils.data import IterableDataset

class LargeIterableDataset(IterableDataset):
    def __init__(self, x_path, y_path):
        self.x_path = x_path
        self.y_path = y_path
        
        with open(x_path + '.conf', 'r') as f:
            x_config = json.load(f)
        self.total_samples = x_config['shape'][0]
        self.x_shape = tuple(x_config['shape'])
        self.x_dtype = x_config['dtype']
        
        with open(y_path + '.conf', 'r') as f:
            y_config = json.load(f)
        self.y_shape = tuple(y_config['shape'])
        self.y_dtype = y_config['dtype']

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        # 分配每个子进程的处理范围
        if worker_info is None:
            start, end = 0, self.total_samples
        else:
            per_worker = self.total_samples // worker_info.num_workers
            start = worker_info.id * per_worker
            end = start + per_worker
            if worker_info.id == worker_info.num_workers - 1:
                end = self.total_samples
        
        # 打开当前进程的MemMap
        x_mem = np.memmap(self.x_path, mode='r', shape=self.x_shape, dtype=self.x_dtype)
        y_mem = np.memmap(self.y_path, mode='r', shape=self.y_shape, dtype=self.y_dtype)
        
        for idx in range(start, end):
            yield x_mem[idx], y_mem[idx]

# 构建DataLoader
def iterable_dataloader(x_path, y_path, batch_size):
    dset = LargeIterableDataset(x_path, y_path)
    return DataLoader(
        dataset=dset,
        batch_size=batch_size,
        num_workers=4,
        pin_memory=True
    )

4. 调整DataLoader与系统参数 #

• 关闭pin_memory：如果你的训练未使用GPU，pin_memory=True会额外占用内存，可改为False；若使用GPU，可尝试限制每个worker的内存占用。
• 修改进程共享策略：设置PyTorch的多进程共享策略为文件系统，避免文件描述符冲突：

  import torch.multiprocessing as mp
  mp.set_sharing_strategy('file_system')
  

• 增大系统Swap空间：若工作站物理内存不足，增大Swap可避免因内存耗尽导致的冻结重启（注意Swap速度较慢，仅作为临时方案）。

5. 限制并发进程数 #

若上述方案仍无法解决，可通过系统工具限制Python进程的内存占用，或调整num_workers为2-3，同时避免同时运行多个脚本，观察是否能稳定运行。

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