streaming_loader.py

import numpy as np
import os
import pickle as pkl
from tqdm import tqdm

"""
Write data into GPFS
"""
NUM_WRITE_PKL = 100
DATA_PATH = "/mnt/work/debug_data/"
NUM_LOOP_LOADING = 100000

for i in tqdm(range(NUM_WRITE_PKL)):
    arr = np.random.random((72, 128, 128, 3))
    with open(os.path.join(DATA_PATH, "test" + str(i).zfill(4) + ".pkl"), "wb") as f:
        pkl.dump(arr, f)
print("Write Files Done")


"""
Read files from GPFS
"""
load_counter = 0
while load_counter < NUM_LOOP_LOADING:
    index = np.random.choice(NUM_WRITE_PKL)
    with open(os.path.join(DATA_PATH, "test" + str(index).zfill(4) + ".pkl"), "rb") as f:
        data = pkl.load(f)
    load_counter += 1
    if (load_counter % 1000) == 0:
        print("Current counter: %i, data shape: %s" % (load_counter, data.shape))
        
print("All done")