Class _BaseDataModule

datamodule.py:42–87

class _BaseDataModule(LightningDataModule):
    def __init__(
        self,
        data_dir: str,
        batch_size: int = 64,
        num_workers: int = 4,
        pin_memory: bool = True,
        size: int = 224,
        augment: bool = True,
        num_samples: Optional[int] = None,
    ):
        super().__init__()

        self.augment = augment
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.pin_memory = pin_memory
        if isinstance(size, int):
            self.size = (size, size)
        else:
            self.size = size
        self.num_samples = num_samples

    def setup(self, stage=None):
        pass

    def train_dataloader(self):
        return DataLoader(
            dataset = self.data_train,
            batch_size = self.batch_size,
            num_workers = self.num_workers,
            pin_memory = self.pin_memory,
            shuffle = True,
            drop_last = True
        )

    def val_dataloader(self):
        return DataLoader(
            dataset = self.data_val,
            batch_size = self.batch_size,
            num_workers = self.num_workers,
            pin_memory = self.pin_memory,
            shuffle = False,
            drop_last = False
        )


class _FGVCDataModule(_BaseDataModule):