LMDB转HDF5（约370万条数据）时写入性能随时间下降的问题排查与优化方案咨询

看起来你遇到的是大规模HDF5写入的典型性能陷阱，我之前处理过千万级别的数据集转换，结合你的代码和场景，咱们一步步拆解问题和优化方案：

为什么会越写越慢？ #

你前150万条快、后期性能暴跌的核心原因，和HDF5的元数据管理方式直接相关：

• 零散子数据集的元数据爆炸：你现在给每个样本单独创建一个子数据集（node_features_group.create_dataset(str(i), ...)），每创建一个子数据集，HDF5都要更新文件的元数据索引。当子数据集数量破百万后，元数据会变得异常庞大，磁盘寻道和索引更新的开销会指数级上升——这就是性能滑坡的罪魁祸首。
• 小数据块压缩的额外开销：你给每个极小的子数据集单独开了gzip压缩，gzip对小数据块的压缩效率极低，还会进一步膨胀元数据，雪上加霜。
• 内存碎片与缓存累积：虽然你加了gc.collect()，但频繁创建销毁小numpy数组，加上HDF5的缓存逐渐累积，还是会导致内存碎片，拖慢后续IO操作。

针对性优化方案（结合你的代码修改） #

1. 核心优化：把零散子数据集改成大型分块数据集（重中之重） #

不要给每个样本单独建子数据集，而是预先定义大型的、支持分块的可变长度数据集，按顺序写入。这样HDF5只需要维护少量元数据，写入性能会全程稳定。

2. 优化压缩与IO策略 #

换用轻量级压缩算法，减少磁盘IO的额外开销；同时手动控制缓存和flush时机，避免频繁磁盘寻道。

3. 内存与循环效率优化 #

预分配数据结构、减少不必要的切片操作，进一步降低内存压力。

修改后的完整优化代码 #

import lmdb
import pickle
import zlib
import numpy as np
import h5py
from tqdm import tqdm
import os
import gc
import time

def convert_pretrain_to_h5( lmdb_path, output_dir, train_keys_path, val_keys_path, test_keys_path ):
    start_time = time.time()
    os.makedirs(output_dir, exist_ok=True)

    # 加载拆分键并保存原始拆分索引
    train_keys = np.load(train_keys_path, allow_pickle=True)
    val_keys = np.load(val_keys_path, allow_pickle=True)
    test_keys = np.load(test_keys_path, allow_pickle=True)

    splits = {
        'train': list(train_keys),
        'val': list(val_keys),
        'test': list(test_keys)
    }
    with open(os.path.join(output_dir, 'split_indices.p'), 'wb') as f:
        pickle.dump(splits, f)
    print(f"Saved original split indices with {len(splits['train'])} train, {len(splits['val'])} val, {len(splits['test'])} test samples")

    all_keys = list(train_keys) + list(val_keys) + list(test_keys)
    total_samples = len(all_keys)
    print(f"Total keys from split files: {len(all_keys)}")

    # 预定义HDF5数据类型（支持可变长度数组，适配不同样本的节点/边数量）
    dt_node = h5py.vlen_dtype(np.dtype('int16'))
    dt_edge = h5py.vlen_dtype(np.dtype('int32'))
    h5_path = os.path.join(output_dir, 'substructure_graphs.h5')

    # 打开LMDB环境（保持原有优化参数）
    env = lmdb.open(lmdb_path, readonly=True, lock=False, readahead=False, meminit=False)
    with env.begin(write=False) as txn:
        # 创建HDF5文件，启用最新版本特性并设置缓存大小
        with h5py.File(h5_path, 'w', libver='latest', swmr=False, rdcc_nbytes=512*1024*1024) as graphs_h5:
            # 预定义大型分块数据集
            node_features_dset = graphs_h5.create_dataset(
                'node_features', shape=(total_samples,), dtype=dt_node,
                chunks=(100000,), compression='lzf'  # lzf压缩速度远快于gzip，适合大规模数据
            )
            edge_index_dset = graphs_h5.create_dataset(
                'edge_index', shape=(total_samples,), dtype=dt_edge,
                chunks=(100000,), compression='lzf'
            )
            num_nodes_dset = graphs_h5.create_dataset('num_nodes', (total_samples,), dtype=np.int32)
            num_edges_dset = graphs_h5.create_dataset('num_edges', (total_samples,), dtype=np.int32)
            smiles_dset = graphs_h5.create_dataset(
                'smiles', shape=(total_samples,), dtype=h5py.string_dtype(encoding='utf-8'),
                chunks=(100000,)
            )

            chunk_size = 100000
            flush_interval = 3  # 每3个chunk手动刷新一次，平衡IO与性能
            chunk_count = 0

            for start in range(0, total_samples, chunk_size):
                end = min(start + chunk_size, total_samples)
                print(f"Processing chunk {start} to {end}...")
                chunk_count += 1

                # 直接遍历全局索引，避免切片开销
                for global_idx in tqdm(range(start, end)):
                    i = all_keys[global_idx]
                    key = f"{i}".encode("ascii")
                    try:
                        data = txn.get(key)
                        sample = pickle.loads(zlib.decompress(data))

                        # 写入SMILES
                        if 'smiles' in sample:
                            smiles = sample['smiles']
                            smiles_dset[global_idx] = smiles.decode('utf-8') if isinstance(smiles, bytes) else smiles

                        # 写入图特征
                        if 'node_features' in sample and 'edge_index' in sample:
                            # 转换并写入节点特征
                            node_feat = np.array(sample['node_features'], dtype=np.int16) if not isinstance(sample['node_features'], np.ndarray) else sample['node_features'].astype(np.int16)
                            node_features_dset[global_idx] = node_feat

                            # 转换并写入边索引
                            edge_idx = np.array(sample['edge_index'], dtype=np.int32) if not isinstance(sample['edge_index'], np.ndarray) else sample['edge_index'].astype(np.int32)
                            edge_index_dset[global_idx] = edge_idx

                        # 写入图元数据
                        if 'num_nodes' in sample:
                            num_nodes_dset[global_idx] = sample['num_nodes']
                        if 'num_edges' in sample:
                            num_edges_dset[global_idx] = sample['num_edges']

                    except Exception as e:
                        print(f"Error processing sample {i}: {e}")
                        raise e

                # 定期刷新缓存并清理内存
                if chunk_count % flush_interval == 0:
                    graphs_h5.flush()
                    gc.collect()
                    print(f"Flushed data to disk, progress: {end}/{total_samples}")

            # 最后一次强制刷新所有数据
            graphs_h5.flush()

    print(f"Conversion complete! Files saved to {output_dir}")
    print(f"Total samples processed: {total_samples}")
    print(f"Total time elapsed: {time.time() - start_time:.2f} seconds")

额外性能小技巧 #

1. 磁盘分离：如果有条件，把LMDB源文件和HDF5输出文件放在不同物理磁盘上，避免读写IO竞争。
2. 压缩取舍：如果磁盘空间充足，可以完全关闭压缩（去掉compression='lzf'），写入速度会再提升30%-50%。
3. chunk_size调整：根据你的内存大小和磁盘带宽，调整chunk_size（比如改成200000），找到最适合你的平衡点。

按照这个方案修改后，你应该能看到全程稳定的写入速度，不会再出现后期性能暴跌的情况。如果还有瓶颈，可以再检查磁盘IO是否达到硬件上限~

内容来源于stack exchange