Function get_pipe_dataset

dali/test/python/test_utils_tensorflow.py:191–224

(
    batch_size,
    num_threads,
    device,
    device_id,
    num_devices=1,
    *,
    dali_on_dev_0=True,
    get_pipeline_desc=get_image_pipeline,
    to_dataset=to_image_dataset,
)

def get_pipe_dataset(
    batch_size,
    num_threads,
    device,
    device_id,
    num_devices=1,
    *,
    dali_on_dev_0=True,
    get_pipeline_desc=get_image_pipeline,
    to_dataset=to_image_dataset,
):
    shard_id = 0 if num_devices == 1 else device_id

    tf_dataset = get_dali_dataset(
        batch_size,
        num_threads,
        device,
        device_id,
        num_devices=num_devices,
        get_pipeline_desc=get_pipeline_desc,
        to_dataset=to_dataset,
    )

    dali_pipeline, _, _ = get_pipeline_desc(
        batch_size,
        num_threads,
        device,
        0 if dali_on_dev_0 else device_id,
        shard_id,
        num_devices,
        def_for_dataset=False,
    )

    return dali_pipeline, tf_dataset


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