File size: 713 Bytes
6325697 |
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 |
from .dataset import Dataset, ValDataset, TestDataset
from torch.utils.data import DataLoader
def find_dataset_using_name(name):
mapping = {
"Video": Dataset,
"VideoVal": ValDataset,
"VideoTest": TestDataset,
}
cls = mapping.get(name, None)
if cls is None:
raise ValueError(f"Fail to find dataset {name}")
return cls
def create_dataset(metainfo, split):
dataset_cls = find_dataset_using_name(split.type)
dataset = dataset_cls(metainfo, split)
return DataLoader(
dataset,
batch_size=split.batch_size,
drop_last=split.drop_last,
shuffle=split.shuffle,
num_workers=split.worker,
pin_memory=True
) |