Function check_prepare_first_batch

dali/test/python/test_fw_iterators.py:1608–1649

(Iterator, *args, to_np=None, **kwargs)

def check_prepare_first_batch(Iterator, *args, to_np=None, **kwargs):
    max_batch_size = 4
    iter_limit = 4
    runs = 3
    test_data_shape = [2, 3, 4]
    i = 0
    dataset = [
        [
            [
                np.random.randint(0, 255, size=test_data_shape, dtype=np.uint8)
                for _ in range(max_batch_size)
            ]
        ]
        for _ in range(iter_limit)
    ]

    def get_data():
        nonlocal i
        if i == iter_limit:
            i = 0
            raise StopIteration
        out = dataset[i]
        i += 1
        return out

    pipe = Pipeline(batch_size=max_batch_size, num_threads=1, device_id=0)
    with pipe:
        outs = fn.external_source(source=get_data, num_outputs=1)
    pipe.set_outputs(*outs)

    it = Iterator([pipe], *args, auto_reset=True, prepare_first_batch=False, **kwargs)
    counter = 0
    for r in range(runs):
        if r == 0:
            # when prepare_first_batch=False pipeline should not be run until first call to next(it)
            assert i == 0, "external_source should not be run yet"
        for j, data in enumerate(it):
            if not isinstance(data, dict):
                data = data[0]
            assert (to_np(data) == np.concatenate(dataset[j])).all()
            counter += 1
    assert counter == iter_limit * runs


@attr("pytorch")