| from header import * |
| from .samplers import DistributedBatchSampler, DistributedMultiDatasetBatchSampler |
| from .catalog import DatasetCatalog |
| from .utils import instantiate_from_config |
| import torch |
| from torch.utils.data import ConcatDataset |
| from .concat_dataset import MyConcatDataset |
|
|
|
|
| def load_dataset(args, dataset_name_list): |
| """ |
| Args: |
| args: |
| dataset_name_list: List[str] |
| repeats: List[int], the training epochs for each dataset |
| |
| """ |
| |
| concat_data = MyConcatDataset(dataset_name_list) |
| world_size = torch.distributed.get_world_size() |
| rank = torch.distributed.get_rank() |
| batch_size = args['world_size'] * args['dschf'].config['train_micro_batch_size_per_gpu'] |
| sampler = torch.utils.data.RandomSampler(concat_data) |
| batch_sampler = DistributedMultiDatasetBatchSampler(dataset=concat_data, |
| sampler=sampler, |
| batch_size=batch_size, |
| drop_last=True, |
| rank=rank, |
| world_size=world_size) |
| iter_ = DataLoader( |
| concat_data, |
| batch_sampler=batch_sampler, |
| num_workers=1, |
| collate_fn=concat_data.collate, |
| pin_memory=True |
| ) |
| return concat_data, iter_, sampler |
|
|
