Function test_dali_proxy_torch_data_loader

dali/test/python/test_dali_proxy.py:91–158

(device, include_decoder, debug=False)

@attr("pytorch")
@params(("cpu", False), ("cpu", True), ("gpu", False), ("gpu", True))
def test_dali_proxy_torch_data_loader(device, include_decoder, debug=False):
    # Shows how DALI proxy is used in practice with a PyTorch data loader

    from nvidia.dali.plugin.pytorch.experimental import proxy as dali_proxy
    import torchvision.datasets as datasets
    from torch.utils import data as torchdata

    batch_size = 4
    num_threads = 3
    device_id = 0
    nworkers = 4
    pipe = image_pipe(
        dali_device=device,
        include_decoder=include_decoder,
        random_pipe=False,
        batch_size=batch_size,
        num_threads=num_threads,
        device_id=device_id,
        prefetch_queue_depth=2 + nworkers,
    )

    pipe_ref = image_pipe(
        dali_device=device,
        include_decoder=include_decoder,
        random_pipe=False,
        batch_size=batch_size,
        num_threads=num_threads,
        device_id=device_id,
        prefetch_queue_depth=1,
    )

    dali_server = dali_proxy.DALIServer(pipe)
    if include_decoder:
        dataset = datasets.ImageFolder(jpeg, transform=dali_server.proxy, loader=read_filepath)
        dataset_ref = datasets.ImageFolder(jpeg, transform=lambda x: x.copy(), loader=read_filepath)
    else:
        dataset = datasets.ImageFolder(jpeg, transform=dali_server.proxy)
        dataset_ref = datasets.ImageFolder(jpeg, transform=lambda x: x.copy())

    loader = dali_proxy.DataLoader(
        dali_server,
        dataset,
        batch_size=batch_size,
        num_workers=nworkers,
        drop_last=True,
    )

    def ref_collate_fn(batch):
        filepaths, labels = zip(*batch)  # Separate the inputs and labels
        # Just return the batch as they are, a list of individual tensors
        return filepaths, labels

    loader_ref = torchdata.dataloader.DataLoader(
        dataset_ref,
        batch_size=batch_size,
        num_workers=1,
        collate_fn=ref_collate_fn,
        shuffle=False,