Function test_dali_proxy_deterministic

dali/test/python/test_dali_proxy.py:287–346

(deterministic, debug=False)

@attr("pytorch")
@params((False,), (True,))
def test_dali_proxy_deterministic(deterministic, debug=False):
    # Shows how DALI proxy can be configured for deterministic results
    from nvidia.dali.plugin.pytorch.experimental import proxy as dali_proxy
    import torchvision.datasets as datasets
    import torch

    # Use a high number of iterations for non-deterministic tests, even though
    # we stop the test once we get different results (usually in the first iteration).
    # For deterministic tests, we check that all runs produce the same results.
    niterations = 3 if deterministic else 10
    num_workers = 4
    seed0 = 123456
    seed1 = 5555464
    seed2 = 775653

    outputs = []
    for i in range(niterations):
        pipe = image_pipe(
            random_pipe=True,
            dali_device="gpu",
            batch_size=1,
            num_threads=1,
            device_id=0,
            seed=seed0,
            prefetch_queue_depth=1,
        )
        outputs_i = []
        torch.manual_seed(seed2)
        with dali_proxy.DALIServer(pipe, deterministic=deterministic) as dali_server:
            dataset = datasets.ImageFolder(jpeg, transform=dali_server.proxy, loader=read_filepath)
            # many workers so that we introduce a lot of variability in the order of arrival
            loader = dali_proxy.DataLoader(
                dali_server,
                dataset,
                batch_size=1,
                num_workers=num_workers,
                shuffle=True,
                worker_init_fn=lambda worker_id: np.random.seed(seed1 + worker_id),
            )
            outputs_i = []
            for _ in range(num_workers):
                for data, _ in loader:
                    outputs_i.append(data.cpu())
                    break
        outputs.append(outputs_i)

        if i > 0:
            if deterministic:
                for k in range(num_workers):
                    assert np.array_equal(outputs[i][k], outputs[0][k])
            else:
                for k in range(num_workers):
                    if not np.array_equal(outputs[i][k], outputs[0][k]):
                        return  # OK

        pipe._shutdown()
        del pipe