如何同时从多个H5文件生成数据？现有代码出现死锁

问题分析 #

你遇到的死锁问题核心原因有两个：

1. 生成器与h5py上下文管理器的冲突：_get_generator中使用with h5py.File(...)管理文件，但生成器的暂停特性会让上下文管理器长期持有文件句柄，在多线程环境下（如PyTorch DataLoader默认多线程），多个文件句柄的并发操作容易引发资源竞争。
2. h5py的线程不安全特性：h5py的文件对象并非线程安全，当多个线程同时通过生成器操作不同H5文件时，会导致线程阻塞死锁。

另外，你的多文件Dataset代码中大概率缺失了self._h5_gen的初始化逻辑（需要初始化为对应文件数量的None列表），这也可能加剧问题。

解决方案 #

方案一：放弃生成器，直接单次读取（简单可靠） #

去掉生成器逻辑，每次__getitem__时直接打开文件读取对应索引的数据，虽然每次打开文件有微小开销，但彻底避免了生成器和多线程的冲突：

import h5py
import numpy as np
from torch.utils.data import Dataset

class H5MultiDataset(Dataset):
    def __init__(self, h5_paths):
        super().__init__()
        self.h5_paths = h5_paths

    def __getitem__(self, index):
        features = [self._read_single_file(index, path_id) for path_id in range(len(self.h5_paths))]
        aligned_t, fusion_t, sensor_t, sensor_t_1, pad_num, radar_t, radar_t_1 = zip(*features)
   
        # 数据后处理逻辑保持不变
        pad_num = np.array(pad_num)
        aligned_t = np.array(aligned_t)
        fusion_t = np.array(fusion_t)
        sensor_t_11 = np.array(sensor_t_1)
        sensor_t_1 = sensor_t_11[...,:-1]
        sensor_ids = np.array(sensor_t)[...,-1]
        sensor_t = np.array(sensor_t)[...,:-1]
        radar_t_1 = np.array(radar_t_1)[...,:-1]
        radar_t = np.array(radar_t)[...,:-1]
        
        return aligned_t, fusion_t, sensor_t, sensor_t_1, pad_num, radar_t, radar_t_1, sensor_ids

    def _read_single_file(self, index, path_id):
        with h5py.File(self.h5_paths[path_id], 'r') as record:
            aligned_t = record['aligned_t'][index]
            fusion_t = record['fusion_t'][index]
            sensor_t = record['sensor_t'][index]
            sensor_t_1 = record['sensor_t_1'][index]
            pad_num = record['pad_num'][index]
            radar_t = record['radar_t'][index]
            radar_t_1 = record['radar_t_1'][index]
            return aligned_t, fusion_t, sensor_t, sensor_t_1, pad_num, radar_t, radar_t_1

方案二：线程本地存储管理文件句柄（性能优化） #

如果想避免频繁打开关闭文件的开销，可以用线程本地存储为每个线程维护独立的文件句柄，既保证线程安全，又能复用文件连接：

import h5py
import numpy as np
import threading
from torch.utils.data import Dataset

class H5MultiDataset(Dataset):
    def __init__(self, h5_paths):
        super().__init__()
        self.h5_paths = h5_paths
        self._thread_local = threading.local()

    def _get_thread_handles(self):
        # 为当前线程创建/获取文件句柄
        if not hasattr(self._thread_local, 'handles'):
            self._thread_local.handles = []
            for path in self.h5_paths:
                self._thread_local.handles.append(h5py.File(path, 'r'))
        return self._thread_local.handles

    def __getitem__(self, index):
        handles = self._get_thread_handles()
        features = []
        for handle in handles:
            # 直接通过句柄读取数据
            aligned_t = handle['aligned_t'][index]
            fusion_t = handle['fusion_t'][index]
            sensor_t = handle['sensor_t'][index]
            sensor_t_1 = handle['sensor_t_1'][index]
            pad_num = handle['pad_num'][index]
            radar_t = handle['radar_t'][index]
            radar_t_1 = handle['radar_t_1'][index]
            features.append((aligned_t, fusion_t, sensor_t, sensor_t_1, pad_num, radar_t, radar_t_1))
        
        # 数据后处理逻辑保持不变
        aligned_t, fusion_t, sensor_t, sensor_t_1, pad_num, radar_t, radar_t_1 = zip(*features)
        pad_num = np.array(pad_num)
        aligned_t = np.array(aligned_t)
        fusion_t = np.array(fusion_t)
        sensor_t_11 = np.array(sensor_t_1)
        sensor_t_1 = sensor_t_11[...,:-1]
        sensor_ids = np.array(sensor_t)[...,-1]
        sensor_t = np.array(sensor_t)[...,:-1]
        radar_t_1 = np.array(radar_t_1)[...,:-1]
        radar_t = np.array(radar_t)[...,:-1]
        
        return aligned_t, fusion_t, sensor_t, sensor_t_1, pad_num, radar_t, radar_t_1, sensor_ids

    def __del__(self):
        # 销毁时关闭当前线程的文件句柄
        if hasattr(self._thread_local, 'handles'):
            for handle in self._thread_local.handles:
                handle.close()

方案三：单线程加载（临时应急） #

如果必须保留原生成器逻辑，可以将DataLoader的num_workers设置为0，强制单线程加载，避免多线程冲突，但会降低数据加载效率：

# 初始化DataLoader时设置
dataloader = DataLoader(dataset, batch_size=32, num_workers=0)

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