Spaces:
Running
Running
File size: 1,360 Bytes
d015578 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import spiga.data.loaders.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
|