42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
from torch.utils.data import DataLoader
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
|
|
from .import alignments as zoo_alignments
|
|
|
|
zoos = [zoo_alignments]
|
|
|
|
|
|
def get_dataset(data_config, pretreat=None, debug=False):
|
|
|
|
for zoo in zoos:
|
|
dataset = zoo.get_dataset(data_config, pretreat=pretreat, debug=debug)
|
|
if dataset is not None:
|
|
return dataset
|
|
raise NotImplementedError('Dataset not available')
|
|
|
|
|
|
def get_dataloader(batch_size, data_config, pretreat=None, sampler_cfg=None, debug=False):
|
|
|
|
dataset = get_dataset(data_config, pretreat=pretreat, debug=debug)
|
|
|
|
if (len(dataset) % batch_size) == 1 and data_config.shuffle == True:
|
|
drop_last_batch = True
|
|
else:
|
|
drop_last_batch = False
|
|
|
|
shuffle = data_config.shuffle
|
|
sampler = None
|
|
if sampler_cfg is not None:
|
|
sampler = DistributedSampler(dataset, num_replicas=sampler_cfg.world_size, rank=sampler_cfg.rank)
|
|
shuffle = False
|
|
|
|
dataloader = DataLoader(dataset,
|
|
batch_size=batch_size,
|
|
shuffle=shuffle,
|
|
num_workers=data_config.num_workers,
|
|
pin_memory=True,
|
|
drop_last=drop_last_batch,
|
|
sampler=sampler)
|
|
|
|
return dataloader, dataset
|