执行其他任务时预加载部分HDF5文件

时间:2015-12-25 17:47:09

标签: python multithreading loading h5py

我正在训练深度学习分类器,它使用的HDF5数据集太大而无法放入内存中。因此,我以256个批量提取数据,并使用这些批次以下列方式训练我的分类器。我使用的深度学习库(Keras)提供了方法model.train_on_batch(X_batch, y_batch)

for i in range(n_batches_in_dset):
        X_batch, y_batch = load_partition('train', ind=[i*batch_size, (i+1)*batch_size])
        loss = model.train_on_batch(X_batch, y_batch)

在使用GPU训练当前数据的同时预取下一批数据是有意义的。如何在Python中做到这一点?

我附上了用于加载数据的代码。

def load_hdf5(path, datapart, ind=None):
    f   = h5py.File(path, 'r')
    if ind is None:
        dat = f[datapart][:]
    else:
        dat = f[datapart][ind[0]:ind[1]]
    f.close()
    return np.array(dat)

def load_partition(name, ind=None):
    path = DEEP_ROOT + 'data/{}.h5'.format(name)
    X = load_hdf5(path, 'data', ind)
    y = load_hdf5(path, 'label', ind)
    X = np.swapaxes(X, 2, 3)
    y = np_utils.to_categorical(y)
    return X, y

1 个答案:

答案 0 :(得分:4)

可能最简单的事情就是将单独的任务放在单独的threads中,并使用同步queue来交换它们之间的批次。我们将为数据读取部分使用单独的线程,并为训练部分使用主线程。

import Queue, threading

data_queue = Queue.Queue(2) # a queue with two space for two "chunks"
sentinel = object()

#start the data-loading task
def load_task()
    for x in i in range(n_batches_in_dset):
        data_queue.put(load_partition('train', ind=[i*batch_size, (i+1)*batch_size]), True)
    # tell the other side we're "done"
    data_queue.put(sentinel, True)

threading.Thread(target=load_task).start()

while True:
    batch = data_queue.get(True)
    data_queue.task_done()
    if batch is sentinel:
        break # we're done now!
    X_batch, y_batch = batch
    loss = model.train_on_batch(X_batch, y_batch)

编辑:我们需要使用Queue.task_done()取消阻止队列