def collect_dataset_dicts(d2_dataset, num_workers_per_gpu=8, dummy_batch_size=32):
    """Build D2 dataset (i.e. List[Dict]), given a dataset implementing recipe for building its item.

    This is useful when a __getitem__() takes much time / memory, so it's desirable to do one-time
    conversion into `List[Dict]` format. This function repurpose multi-node distributed training tools
    of pytorch to acccelerate data loading.

    NOTE: Both the input dataset object and the resulting D2 dataset, multiplied by # of GPUs, must fit in memory.

    Parameters
    ----------
    d2_dataset: Dataset
        __getitem__() returns a D2-formatted dictionary.
    """
    class TrivialModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.theta = nn.Parameter(torch.FloatTensor([0.]))

        def forward(self, x):  # pylint: disable=unused-argument
            return self.theta * 0.

    model = TrivialModel().to('cuda')

    if is_distributed():
        model = DistributedDataParallel(
            model,
            device_ids=[comm.get_local_rank()],
            broadcast_buffers=False,
        )
        sampler = InferenceSampler(len(d2_dataset))  # d2 dependency
    else:
        sampler = None

    dataloader = torch.utils.data.DataLoader(
        d2_dataset,
        num_workers=num_workers_per_gpu,
        sampler=sampler,
        batch_size=dummy_batch_size,
        collate_fn=lambda x: x,
        drop_last=False,
    )
    d2_dicts = []
    LOG.info(f"Creating D2 dataset: {len(dataloader)} batches on rank {comm.get_rank()}.")
    for x in tqdm(dataloader):
        loss = model(x)
        loss.backward()
        d2_dicts.extend(x)

    LOG.info("Gathering D2 dataset dicts from all GPU workers...")
    st = time.time()
    d2_dicts = MPI.COMM_WORLD.allgather(d2_dicts)
    took = time.time() - st
    d2_dicts = list(itertools.chain.from_iterable(d2_dicts))
    LOG.info(f"Done (length={len(d2_dicts)}, took={took:.1f}s).")

    return d2_dicts
Build D2 dataset (i.e. List[Dict]), given a dataset implementing recipe for building its item.

    This is useful when a __getitem__() takes much time / memory, so it's desirable to do one-time
    conversion into `List[Dict]` format. This function repurpose multi-node distributed training tools
    of pytorch to acccelerate data loading.

    NOTE: Both the input dataset object and the resulting D2 dataset, multiplied by # of GPUs, must fit in memory.

    Parameters
    ----------
    d2_dataset: Dataset
        __getitem__() returns a D2-formatted dictionary.
    
tridet/data/build.py