Function build_data_loader

linear_moe/finetune_utils.py:82–105

Data loader. Note that batch-size is the local (per GPU) batch-size.

(dataset,
                      micro_batch_size,
                      num_workers,
                      drop_last,
                      task_collate_fn=None)

def build_data_loader(dataset,
                      micro_batch_size,
                      num_workers,
                      drop_last,
                      task_collate_fn=None):
    """Data loader. Note that batch-size is the local (per GPU) batch-size."""

    # Sampler.
    world_size = mpu.get_data_parallel_world_size()
    rank = mpu.get_data_parallel_rank()
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank)

    # Data loader. Note that batch size is the per GPU batch size.
    data_loader = torch.utils.data.DataLoader(dataset,
                                              batch_size=micro_batch_size,
                                              sampler=sampler,
                                              shuffle=False,
                                              num_workers=num_workers,
                                              drop_last=drop_last,
                                              pin_memory=True,
                                              collate_fn=task_collate_fn)

    return data_loader


def _build_infinite_size_dataloader(dataloader):