[Python] HDF5 Write Performance Degrades Over Time When Converting from LMDB (~3.7M entries)

I’m experiencing significant slow-downs when converting data from LMDB to HDF5 format. While the conversion starts off quickly, performance degrades substantially partway through the process.

Specifically, my dataset contains around 3.7 million entries, and the first ~1.5 million entries are processed quite fast. However, after that point, the process becomes increasingly slow.

I’m already processing the data in chunks to manage memory and improve performance, but this hasn’t solved the issue.

I’m suspecting the issue might relate to how I’m creating the HDF5 file. Are there any optimizations I can apply to reduce this performance degradation during dataset creation?

Any suggestions for improving performance — or explanations for why this slowdown happens — would be greatly appreciated.

Please let me know if you’d like me to share additional information.

Thanks in advance!

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)

# Load split keys (preserving original indices)
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)

# Save the original split indices directly - this preserves your exact splits
splits = {
'train': list(train_keys),
'val': list(val_keys),
'test': list(test_keys)
}

# Save split indices with original values
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")

# Now use all_keys with original values for processing
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)}")

# Create H5 files
h5_path = os.path.join(output_dir, 'substructure_graphs.h5')
smiles_list = []

# Get available LMDB keys for verification
env = lmdb.open(lmdb_path, readonly=True, lock=False, readahead=False, meminit=False)

# Create H5 file
with h5py.File(h5_path, 'w') as graphs_h5:
# Create groups for graph data
node_features_group = graphs_h5.create_group('node_features')
edge_index_group = graphs_h5.create_group('edge_index')
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)

# Process in chunks for memory efficiency
chunk_size = 100000
with env.begin(write=False) as txn:
for start in range(0, total_samples, chunk_size):
end = min(start + chunk_size, total_samples)
print(f"Processing chunk {start} to {end}...")

# Process each sample in this chunk
for idx, i in enumerate(tqdm(all_keys[start:end])):
key = f"{i}".encode("ascii")
try:
# Get data from LMDB
data = txn.get(key)

# Decompress and load
sample = pickle.loads(zlib.decompress(data))

# Extract SMILES
if 'smiles' in sample:
smiles = sample['smiles']
if isinstance(smiles, bytes):
smiles = smiles.decode('utf-8')
smiles_list.append(smiles)

# Store graph data
if 'node_features' in sample and 'edge_index' in sample:
# Node features
node_features = sample['node_features']
if not isinstance(node_features, np.ndarray):
node_features = np.array(node_features, dtype=np.int16)

node_features_group.create_dataset(
str(i),
data=node_features,
compression="gzip" # Faster without compression
)

# Edge indices
edge_index = sample['edge_index']
if not isinstance(edge_index, np.ndarray):
edge_index = np.array(edge_index, dtype=np.int32)

edge_index_group.create_dataset(
str(i),
data=edge_index,
compression="gzip" # Faster without compression
)

# Graph metadata
if 'num_nodes' in sample:
num_nodes_dset = sample['num_nodes']
if 'num_edges' in sample:
num_edges_dset = sample['num_edges']

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


# Save SMILES after each chunk for safety
np.save(os.path.join(output_dir, 'smiles.npy'), np.array(smiles_list, dtype=object))
print(f"Saved progress: {len(smiles_list)}/{total_samples} samples")

# After processing a chunk (just before next chunk starts)
graphs_h5.flush() # Force HDF5 to write buffered data to disk
gc.collect() # Clear unused memory (buffers, zlib, arrays)

# Final save of SMILES
np.save(os.path.join(output_dir, 'smiles.npy'), np.array(smiles_list, dtype=object))

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

Continue reading...