diff --git a/python/hdf5.py b/python/hdf5.py index 4acbccfb3acca59fbff54e052f831442218d156b..13681011d82fa99e7a21f4e533dbc75592aa4521 100644 --- a/python/hdf5.py +++ b/python/hdf5.py @@ -142,20 +142,19 @@ if __name__ == '__main__': #%% import h5py -def hdf5Generator(filePath, batch_size, dataSet): +def hdf5Generator(filePath, batch_size, dataSet,loop=True): with h5py.File(filePath, 'r') as hf: L = len(hf["X" + dataSet]) while True: batch_start = 0 batch_end = batch_size - while batch_start < L: - limit = min(batch_end, L) - X = hf["X" + dataSet][batch_start:limit] - Y = hf["Y" + dataSet][batch_start:limit] - + while batch_end < L: + X = hf["X" + dataSet][batch_start:batch_end] + Y = hf["Y" + dataSet][batch_start:batch_end] yield (X,Y) #a tuple with two numpy arrays with batch_size samples batch_start += batch_size batch_end += batch_size + if not loop: break \ No newline at end of file