batch_size: 256 # batch size for trainset
val_batch_size: 512 # batch size for valset in DatasetInterface

# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE

train_dataset_params:
  root: /data/cifar100
  train: True
  transforms:
    - RandomCrop:
        size: 32
        padding: 4
    - RandomHorizontalFlip
    - ToTensor
    - Normalize:
        mean:
          - 0.5071
          - 0.4865
          - 0.4409
        std:
          - 0.2673
          - 0.2564
          - 0.2762
  target_transform: null
  download: True

train_dataloader_params:
  batch_size: 256
  num_workers: 8
  drop_last: False
  pin_memory: True

val_dataset_params:
  root: /data/cifar100
  train: False
  transforms:
    - ToTensor
    - Normalize:
        mean:
          - 0.5071
          - 0.4865
          - 0.4409
        std:
          - 0.2673
          - 0.2564
          - 0.2762
  target_transform: null
  download: True

val_dataloader_params:
  batch_size: 512
  num_workers: 8
  drop_last: False
  pin_memory: True

@@ -1,2 +1,54 @@
 batch_size: 256 # batch size for trainset
-val_batch_size: 512 # batch size for valset in DatasetInterface
+val_batch_size: 512 # batch size for valset in DatasetInterface
+
+# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
+
+train_dataset_params:
+  root: /data/cifar10
+  train: True
+  transforms:
+    - RandomCrop:
+        size: 32
+        padding: 4
+    - RandomHorizontalFlip
+    - ToTensor
+    - Normalize:
+        mean:
+          - 0.4914
+          - 0.4822
+          - 0.4465
+        std:
+          - 0.2023
+          - 0.1994
+          - 0.2010
+  target_transform: null
+  download: True
+
+train_dataloader_params:
+  batch_size: 256
+  num_workers: 8
+  drop_last: False
+  pin_memory: True
+
+val_dataset_params:
+  root: /data/cifar10
+  train: False
+  transforms:
+    - ToTensor
+    - Normalize:
+        mean:
+          - 0.4914
+          - 0.4822
+          - 0.4465
+        std:
+          - 0.2023
+          - 0.1994
+          - 0.2010
+  target_transform: null
+  download: True
+
+val_dataloader_params:
+  batch_size: 512
+  num_workers: 8
+  drop_last: False
+  pin_memory: True

@@ -17,6 +17,7 @@ from super_gradients.training.datasets.detection_datasets.pascal_voc_detection i
 from super_gradients.training.utils import get_param
 from super_gradients.training.datasets import ImageNetDataset
 from super_gradients.training.datasets.detection_datasets import COCODetectionDataset
+from super_gradients.training.datasets.classification_datasets.cifar import Cifar10, Cifar100
 from super_gradients.training.datasets.segmentation_datasets import CityscapesDataset, CoCoSegmentationDataSet, \
     PascalAUG2012SegmentationDataSet, \
     PascalVOC2012SegmentationDataSet, SuperviselyPersonsDataset
@@ -237,6 +238,42 @@ def tiny_imagenet_val(dataset_params={}, dataloader_params={}, config_name="tiny
                            dataloader_params=dataloader_params)
 
 
+def cifar10_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+    return get_data_loader(config_name="cifar10_dataset_params",
+                           dataset_cls=Cifar10,
+                           train=True,
+                           dataset_params=dataset_params,
+                           dataloader_params=dataloader_params
+                           )
+
+
+def cifar10_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+    return get_data_loader(config_name="cifar10_dataset_params",
+                           dataset_cls=Cifar10,
+                           train=False,
+                           dataset_params=dataset_params,
+                           dataloader_params=dataloader_params
+                           )
+
+
+def cifar100_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+    return get_data_loader(config_name="cifar100_dataset_params",
+                           dataset_cls=Cifar100,
+                           train=True,
+                           dataset_params=dataset_params,
+                           dataloader_params=dataloader_params
+                           )
+
+
+def cifar100_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+    return get_data_loader(config_name="cifar100_dataset_params",
+                           dataset_cls=Cifar100,
+                           train=False,
+                           dataset_params=dataset_params,
+                           dataloader_params=dataloader_params
+                           )
+
+
 def classification_test_dataloader(batch_size: int = 5, image_size: int = 32) -> DataLoader:
     images = torch.Tensor(np.zeros((batch_size, 3, image_size, image_size)))
     ground_truth = torch.LongTensor(np.zeros((batch_size)))

@@ -5,7 +5,7 @@ from super_gradients.training.datasets.data_augmentation import DataAugmentation
 from super_gradients.training.datasets.sg_dataset import ListDataset, DirectoryDataSet
 from super_gradients.training.datasets.all_datasets import CLASSIFICATION_DATASETS, OBJECT_DETECTION_DATASETS, \
     SEMANTIC_SEGMENTATION_DATASETS
-from super_gradients.training.datasets.classification_datasets import ImageNetDataset
+from super_gradients.training.datasets.classification_datasets import ImageNetDataset, Cifar10, Cifar100
 from super_gradients.training.datasets.detection_datasets import DetectionDataset, COCODetectionDataset, PascalVOCDetectionDataset
 from super_gradients.training.datasets.segmentation_datasets.segmentation_dataset import SegmentationDataSet
 from super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmentation import PascalVOC2012SegmentationDataSet
@@ -31,4 +31,5 @@ __all__ = ['DataAugmentation', 'ListDataset', 'DirectoryDataSet', 'CLASSIFICATIO
            'TestYoloDetectionDatasetInterface', 'DetectionTestDatasetInterface', 'ClassificationTestDatasetInterface',
            'SegmentationTestDatasetInterface',
            'ImageNetDatasetInterface',
-           'DetectionDataset', 'COCODetectionDataset', 'PascalVOCDetectionDataset', 'ImageNetDataset', 'SuperviselyPersonsDataset']
+           'DetectionDataset', 'COCODetectionDataset', 'PascalVOCDetectionDataset', 'ImageNetDataset',
+           'Cifar10', 'Cifar100', 'SuperviselyPersonsDataset']

@@ -1,4 +1,5 @@
 from super_gradients.training.datasets.classification_datasets.imagenet_dataset import ImageNetDataset
+from super_gradients.training.datasets.classification_datasets.cifar import Cifar10, Cifar100
 
 
-__all__ = ['ImageNetDataset']
+__all__ = ['ImageNetDataset', 'Cifar10', 'Cifar100']

from typing import Optional, Callable

from torchvision.transforms import Compose

from super_gradients.common.factories.list_factory import ListFactory
from super_gradients.common.factories.transforms_factory import TransformsFactory
from super_gradients.common.decorators.factory_decorator import resolve_param
from torchvision.datasets import CIFAR10, CIFAR100


class Cifar10(CIFAR10):
    """
    CIFAR10 Dataset

    :param root:                    Path for the data to be extracted
    :param train:                   Bool to load training (True) or validation (False) part of the dataset
    :param transforms:              List of transforms to apply sequentially on sample. Wrapped internally with torchvision.Compose
    :param target_transform:        Transform to apply to target output
    :param download:                Download (True) the dataset from source
    """
    @resolve_param("transforms", ListFactory(TransformsFactory()))
    def __init__(
        self,
        root: str,
        train: bool = True,
        transforms: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:
        super(Cifar10, self).__init__(
            root=root,
            train=train,
            transform=Compose(transforms),
            target_transform=target_transform,
            download=download,
        )


class Cifar100(CIFAR100):
    @resolve_param("transforms", ListFactory(TransformsFactory()))
    def __init__(
        self,
        root: str,
        train: bool = True,
        transforms: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:
        """
        CIFAR100 Dataset

        :param root:                    Path for the data to be extracted
        :param train:                   Bool to load training (True) or validation (False) part of the dataset
        :param transforms:              List of transforms to apply sequentially on sample. Wrapped internally with torchvision.Compose
        :param target_transform:        Transform to apply to target output
        :param download:                Download (True) the dataset from source
        """
        super(Cifar100, self).__init__(
            root=root,
            train=train,
            transform=Compose(transforms),
            target_transform=target_transform,
            download=download,
        )

@@ -2,6 +2,6 @@
 
 from tests.end_to_end_tests.trainer_test import TestTrainer
 
-from tests.end_to_end_tests.cifar10_trainer_test import TestCifar10Trainer
+from tests.end_to_end_tests.cifar_trainer_test import TestCifarTrainer
 
-__all__ = ['TestTrainer', 'TestCifar10Trainer']
+__all__ = ['TestTrainer', 'TestCifarTrainer']

import unittest

from super_gradients.training import models

import super_gradients

from super_gradients import Trainer
from super_gradients.training.datasets.dataset_interfaces import LibraryDatasetInterface


class TestCifar10Trainer(unittest.TestCase):
    def test_train_cifar10(self):
        super_gradients.init_trainer()
        trainer = Trainer("test", model_checkpoints_location='local')
        cifar_10_dataset_interface = LibraryDatasetInterface(name="cifar10")
        trainer.connect_dataset_interface(cifar_10_dataset_interface)
        model = models.get("resnet18_cifar", arch_params={"num_classes": 10})
        trainer.train(model=model, training_params={"max_epochs": 1})


if __name__ == '__main__':
    unittest.main()

import unittest

from super_gradients.training import models

import super_gradients

from super_gradients import Trainer
from super_gradients.training.datasets.dataset_interfaces import LibraryDatasetInterface
from super_gradients.training.dataloaders.dataloader_factory import (
    cifar10_train,
    cifar10_val,
    cifar100_train,
    cifar100_val,
)


class TestCifarTrainer(unittest.TestCase):
    def test_train_cifar10(self):
        super_gradients.init_trainer()
        trainer = Trainer("test", model_checkpoints_location="local")
        cifar_10_dataset_interface = LibraryDatasetInterface(name="cifar10")
        trainer.connect_dataset_interface(cifar_10_dataset_interface)
        model = models.get("resnet18_cifar", arch_params={"num_classes": 10})
        trainer.train(
            model=model,
            training_params={
                "max_epochs": 1,
                "initial_lr": 0.1,
                "loss": "cross_entropy",
                "train_metrics_list": ["Accuracy"],
                "valid_metrics_list": ["Accuracy"],
                "metric_to_watch": "Accuracy",
            },
        )

    def test_train_cifar10_dataloader(self):
        super_gradients.init_trainer()
        trainer = Trainer("test", model_checkpoints_location="local")
        cifar10_train_dl, cifar10_val_dl = cifar10_train(), cifar10_val()
        model = models.get("resnet18_cifar", arch_params={"num_classes": 10})
        trainer.train(
            model=model,
            training_params={
                "max_epochs": 1,
                "initial_lr": 0.1,
                "loss": "cross_entropy",
                "train_metrics_list": ["Accuracy"],
                "valid_metrics_list": ["Accuracy"],
                "metric_to_watch": "Accuracy",
            },
            train_loader=cifar10_train_dl,
            valid_loader=cifar10_val_dl,
        )

    def test_train_cifar100(self):
        super_gradients.init_trainer()
        trainer = Trainer("test", model_checkpoints_location="local")
        cifar_10_dataset_interface = LibraryDatasetInterface(name="cifar100")
        trainer.connect_dataset_interface(cifar_10_dataset_interface)
        model = models.get("resnet18_cifar", arch_params={"num_classes": 100})
        trainer.train(
            model=model,
            training_params={
                "max_epochs": 1,
                "initial_lr": 0.1,
                "loss": "cross_entropy",
                "train_metrics_list": ["Accuracy"],
                "valid_metrics_list": ["Accuracy"],
                "metric_to_watch": "Accuracy",
            },
        )

    def test_train_cifar100_dataloader(self):
        super_gradients.init_trainer()
        trainer = Trainer("test", model_checkpoints_location="local")
        cifar100_train_dl, cifar100_val_dl = cifar100_train(), cifar100_val()
        model = models.get("resnet18_cifar", arch_params={"num_classes": 100})
        trainer.train(
            model=model,
            training_params={
                "max_epochs": 1,
                "initial_lr": 0.1,
                "loss": "cross_entropy",
                "train_metrics_list": ["Accuracy"],
                "valid_metrics_list": ["Accuracy"],
                "metric_to_watch": "Accuracy",
            },
            train_loader=cifar100_train_dl,
            valid_loader=cifar100_val_dl,
        )


if __name__ == "__main__":
    unittest.main()

@@ -6,6 +6,10 @@ from super_gradients.training.dataloaders.dataloader_factory import (
     classification_test_dataloader,
     detection_test_dataloader,
     segmentation_test_dataloader,
+    cifar10_val,
+    cifar10_train,
+    cifar100_val,
+    cifar100_train,
     coco2017_train,
     coco2017_val,
     coco2017_train_ssd_lite_mobilenet_v2,
@@ -33,15 +37,23 @@ from super_gradients.training.dataloaders.dataloader_factory import (
     pascal_voc_segmentation_train,
     pascal_voc_segmentation_val,
     supervisely_persons_train,
-    supervisely_persons_val, pascal_voc_detection_train, pascal_voc_detection_val
+    supervisely_persons_val,
+    pascal_voc_detection_train,
+    pascal_voc_detection_val,
+)
+from super_gradients.training.datasets import (
+    COCODetectionDataset,
+    ImageNetDataset,
+    PascalAUG2012SegmentationDataSet,
+    PascalVOC2012SegmentationDataSet,
+    SuperviselyPersonsDataset,
+    PascalVOCDetectionDataset,
+    Cifar10,
+    Cifar100,
 )
-from super_gradients.training.datasets import COCODetectionDataset, ImageNetDataset, PascalAUG2012SegmentationDataSet, \
-    PascalVOC2012SegmentationDataSet, \
-    SuperviselyPersonsDataset, PascalVOCDetectionDataset
 
 
 class DataLoaderFactoryTest(unittest.TestCase):
-
     def test_coco2017_train_creation(self):
         dl_train = coco2017_train()
         self.assertTrue(isinstance(dl_train, DataLoader))
@@ -124,7 +136,9 @@ class DataLoaderFactoryTest(unittest.TestCase):
 
     def test_imagenet_resnet50_kd_train_creation(self):
         # Here we need to overwrite the sampler because the RepeatAugSampler used in KD is only supported for DDP
-        dl = imagenet_resnet50_kd_train(dataloader_params={"sampler": {"InfiniteSampler": {}}})
+        dl = imagenet_resnet50_kd_train(
+            dataloader_params={"sampler": {"InfiniteSampler": {}}}
+        )
         self.assertTrue(isinstance(dl, DataLoader))
         self.assertTrue(isinstance(dl.dataset, ImageNetDataset))
 
@@ -153,6 +167,26 @@ class DataLoaderFactoryTest(unittest.TestCase):
         self.assertTrue(isinstance(dl, DataLoader))
         self.assertTrue(isinstance(dl.dataset, ImageNetDataset))
 
+    def test_cifar10_train_creation(self):
+        dl_train = cifar10_train()
+        self.assertTrue(isinstance(dl_train, DataLoader))
+        self.assertTrue(isinstance(dl_train.dataset, Cifar10))
+
+    def test_cifar10_val_creation(self):
+        dl_val = cifar10_val()
+        self.assertTrue(isinstance(dl_val, DataLoader))
+        self.assertTrue(isinstance(dl_val.dataset, Cifar10))
+
+    def test_cifar100_train_creation(self):
+        dl_train = cifar100_train()
+        self.assertTrue(isinstance(dl_train, DataLoader))
+        self.assertTrue(isinstance(dl_train.dataset, Cifar100))
+
+    def test_cifar100_val_creation(self):
+        dl_val = cifar100_val()
+        self.assertTrue(isinstance(dl_val, DataLoader))
+        self.assertTrue(isinstance(dl_val.dataset, Cifar100))
+
     def test_classification_test_dataloader_creation(self):
         dl = classification_test_dataloader()
         self.assertTrue(isinstance(dl, DataLoader))
@@ -209,5 +243,5 @@ class DataLoaderFactoryTest(unittest.TestCase):
         self.assertTrue(isinstance(dl.dataset, PascalVOCDetectionDataset))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()