|
from abc import ABCMeta, abstractmethod |
|
from functools import lru_cache |
|
from torch.utils.data import DataLoader |
|
|
|
|
|
class BaseDataset(metaclass=ABCMeta): |
|
""" |
|
base class for datasets, it includes 3 types: |
|
- for self-supervised training, |
|
- for classifier training for evaluation, |
|
- for testing |
|
""" |
|
|
|
def __init__( |
|
self, bs_train, aug_cfg, num_workers, bs_clf=1000, bs_test=1000, |
|
): |
|
self.aug_cfg = aug_cfg |
|
self.bs_train, self.bs_clf, self.bs_test = bs_train, bs_clf, bs_test |
|
self.num_workers = num_workers |
|
|
|
@abstractmethod |
|
def ds_train(self): |
|
raise NotImplementedError |
|
|
|
@abstractmethod |
|
def ds_clf(self): |
|
raise NotImplementedError |
|
|
|
@abstractmethod |
|
def ds_test(self): |
|
raise NotImplementedError |
|
|
|
@property |
|
@lru_cache() |
|
def train(self): |
|
return DataLoader( |
|
dataset=self.ds_train(), |
|
batch_size=self.bs_train, |
|
shuffle=True, |
|
num_workers=self.num_workers, |
|
pin_memory=True, |
|
drop_last=True, |
|
) |
|
|
|
@property |
|
@lru_cache() |
|
def clf(self): |
|
return DataLoader( |
|
dataset=self.ds_clf(), |
|
batch_size=self.bs_clf, |
|
shuffle=True, |
|
num_workers=self.num_workers, |
|
pin_memory=True, |
|
drop_last=True, |
|
) |
|
|
|
@property |
|
@lru_cache() |
|
def test(self): |
|
return DataLoader( |
|
dataset=self.ds_test(), |
|
batch_size=self.bs_test, |
|
shuffle=False, |
|
num_workers=self.num_workers, |
|
pin_memory=True, |
|
drop_last=False, |
|
) |
|
|
