Function run_dataset_in_graph

dali/test/python/test_utils_tensorflow.py:227–245

(dali_datasets, iterations, to_stop_iter=False)

def run_dataset_in_graph(dali_datasets, iterations, to_stop_iter=False):
    if not isinstance(dali_datasets, list):
        dali_datasets = [dali_datasets]

    dataset_results = []
    initializers = [tf.compat.v1.global_variables_initializer()]
    ops_to_run = []

    for dali_dataset in dali_datasets:
        iterator = tf.compat.v1.data.make_initializable_iterator(dali_dataset)
        initializers.append(iterator.initializer)
        ops_to_run.append(iterator.get_next())

    with tf.compat.v1.Session() as sess:
        sess.run(initializers)
        with expect_iter_end(not to_stop_iter, tf.errors.OutOfRangeError):
            for _ in range(iterations):
                dataset_results.append(sess.run(ops_to_run))
    return dataset_results


def run_dataset_eager_mode(dali_datasets, iterations, to_stop_iter=False):