如何有效地设置将包含未知数据量的HDF5文件

Question

我有一个可以运行任意长时间的模拟。为了存储模拟的输出，我天真地创建了一个可调整大小的HDF5文件，在获得数据时不断地将数据存储到该文件中，如这个玩具示例所示：

import contextlib
import os
import time
import numpy as np
import h5py
num_timepoints = 18000
num_vertices = 16
num_info = 38
output_size = 10
t0 = "A:\t0.hdf5"
with contextlib.suppress(FileNotFoundError):
os.remove(t0)
st = time.time()
with h5py.File(t0, "a") as f:
dset = f.create_dataset("test", (0, num_vertices, num_info), maxshape=(None, num_vertices, num_info))
for n in np.arange(18000/output_size):
chunk = np.random.rand(output_size, 16, 38)
with h5py.File(t0, "a") as f:
dset = f["test"]
orig_index = dset.shape[0]
dset.resize(dset.shape[0] + chunk.shape[0], axis=0)
dset[orig_index:, :, :] = chunk
et = time.time()
print("test0: time taken: {} s, size: {} kB".format(np.round(et - st, 2), int(os.path.getsize(t0))/1000))

请注意，测试数据的大小与我从模拟中获得的数据的大小平均相似(在最坏的情况下，我可能有测试中时间点数量的2到3倍(。

该测试的输出为：

test0: time taken: 2.02 s, size: 46332.856 kB

将此输出与预先提供数据大小的测试进行比较：

t1 = "A:\t1.hdf5"
with contextlib.suppress(FileNotFoundError):
os.remove(t1)
st = time.time()
data = np.random.rand(num_timepoints, num_vertices, num_info)
with h5py.File(t1, "a") as f:
dset = f.create_dataset("test", data.shape)
dset = data
et = time.time()
print("test1: time taken: {} s, size: {} kB".format(np.round(et - st, 2), int(os.path.getsize(t1))/1000))

其输出为：

test1: time taken: 0.09 s, size: 1.4 kB

如果我选择output_size(它反映了我一次从模拟中获得的数据块有多大(作为1，那么test0大约需要40秒，并创建一个大约700 MB的文件！

显然，test0使用的是一种非常幼稚且效率低下的方法。我该如何改进？我的完整测试代码是：

import contextlib
import os
import time
import numpy as np
import h5py
# =================================================
num_timepoints = 18000
num_vertices = 16
num_info = 38
output_size = 10
t0 = "A:\t0.hdf5"
with contextlib.suppress(FileNotFoundError):
os.remove(t0)
st = time.time()
with h5py.File(t0, "a") as f:
dset = f.create_dataset("test", (0, num_vertices, num_info), maxshape=(None, num_vertices, num_info))
for n in np.arange(18000/output_size):
chunk = np.random.rand(output_size, 16, 38)
with h5py.File(t0, "a") as f:
dset = f["test"]
orig_index = dset.shape[0]
dset.resize(dset.shape[0] + chunk.shape[0], axis=0)
dset[orig_index:, :, :] = chunk
et = time.time()
print("test0: time taken: {} s, size: {} kB".format(np.round(et - st, 2), int(os.path.getsize(t0))/1000))
# =================================================
t1 = "A:\t1.hdf5"
with contextlib.suppress(FileNotFoundError):
os.remove(t1)
st = time.time()
data = np.random.rand(num_timepoints, num_vertices, num_info)
with h5py.File(t1, "a") as f:
dset = f.create_dataset("test", data.shape)
dset = data
et = time.time()
print("test1: time taken: {} s, size: {} kB".format(np.round(et - st, 2), int(os.path.getsize(t1))/1000))
# =================================================
print("Done.")

Answer 1

以下是我发现的一些可以轻松提高性能的东西。首先，不要关闭并重新打开文件来写入每个区块：

with h5py.File(t0, "a") as f:
dset = f["test"]
for n in np.arange(18000/output_size):
chunk = np.random.rand(output_size, 16, 38)
orig_index = dset.shape[0]
dset.resize(dset.shape[0] + chunk.shape[0], axis=0)
dset[orig_index:, :, :] = chunk

这需要大约2秒到大约0.9秒的时间。

其次，h5py为您的数据集猜测了一个相当奇怪的块形状(当我尝试时，是128*4*10(。您可以手动指定要添加的块的形状：

with h5py.File(t0, "a") as f:
dset = f.create_dataset("test", (0, num_vertices, num_info),
maxshape=(None, num_vertices, num_info),
chunks=(output_size, num_vertices, num_info),
)

在这个例子中，我没有得到太多的加速(可能是0.9秒到0.8秒(。但值得一看；它可能会产生更大的差异，这取决于您的数据形状和存储。

最后，如果我一次写入一个更大的块(output_size = 100(，我会看到性能与一次写入的示例相同(或更好(，大约0.5秒(一旦一次写入示例固定为实际写入数据，请参阅我的评论(。

当然，你不想仅仅为了让写作更快而改变你的模拟。但是，如果这种加速很重要，您可以编写一些代码来批量处理模拟中的数据，并定期向HDF5写入更大的块。缺点是，如果模拟崩溃，可能会丢失一些数据。

您也可以减少按较大块调整大小的次数(例如，调整大小以添加100，然后在再次调整大小之前对每行10行进行10次写入(编辑：我试过了，但实际上似乎并没有改善时间。