print("Keras: flow_from_directory")

start_mem = psutil.virtual_memory().used
start_time = time.time()

datagen = ImageDataGenerator(rescale=1./255)
train_generator = datagen.flow_from_directory('flowers', target_size=(128,128), batch_size=32)
print("PRE-TIME", time.time() - start_time)

start_time = time.time()

step = 0
for x_batch, y_batch in train_generator:
    step += 1
    if step > STEPS: break

print("TIME", time.time() - start_time)

curr_mem = psutil.virtual_memory().used
print("Memory Used: %.2f GB" % ((curr_mem - start_mem) / GB))

# Release unused memory

dataset = None
gc.collect()
time.sleep(5)