File size: 920 Bytes
33569f9 | 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 | import os
import torch
from .data_utils import trivial_batch_collator, worker_init_reset_seed
datasets = {}
def register_dataset(name):
def decorator(cls):
datasets[name] = cls
return cls
return decorator
def make_dataset(name, is_training, split, **kwargs):
"""
A simple dataset builder
"""
dataset = datasets[name](is_training, split, **kwargs)
return dataset
def make_data_loader(dataset, is_training, generator, batch_size, num_workers):
"""
A simple dataloder builder
"""
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=trivial_batch_collator,
worker_init_fn=(worker_init_reset_seed if is_training else None),
shuffle=is_training,
drop_last=is_training,
generator=generator,
persistent_workers=True
)
return loader
|