@@ -1,5 +1,4 @@
-from super_gradients.training import ARCHITECTURES, losses, utils, datasets_utils, DataAugmentation, \
-    TestDatasetInterface, SegmentationTestDatasetInterface, DetectionTestDatasetInterface, ClassificationTestDatasetInterface, SgModel, KDModel, \
+from super_gradients.training import ARCHITECTURES, losses, utils, datasets_utils, DataAugmentation, SgModel, KDModel, \
     Trainer, KDTrainer
 from super_gradients.common import init_trainer, is_distributed
 from super_gradients.examples.train_from_recipe_example import train_from_recipe
@@ -7,9 +6,8 @@ from super_gradients.examples.train_from_kd_recipe_example import train_from_kd_
 from super_gradients.sanity_check import env_sanity_check
 
 __all__ = ['ARCHITECTURES', 'losses', 'utils', 'datasets_utils', 'DataAugmentation',
-           'TestDatasetInterface', 'Trainer', 'KDTrainer', 'SegmentationTestDatasetInterface', 'DetectionTestDatasetInterface',
-           'ClassificationTestDatasetInterface', 'init_trainer', 'is_distributed', 'train_from_recipe', 'train_from_kd_recipe',
+           'Trainer', 'KDTrainer',
+           'init_trainer', 'is_distributed', 'train_from_recipe', 'train_from_kd_recipe',
            'env_sanity_check', 'KDModel', 'SgModel']
 
-
 env_sanity_check()

from super_gradients.common.factories.base_factory import BaseFactory
from super_gradients.training.datasets.dataset_interfaces import LibraryDatasetInterface, ClassificationDatasetInterface, Cifar10DatasetInterface,\
    Cifar100DatasetInterface, ImageNetDatasetInterface, TinyImageNetDatasetInterface, CoCoSegmentationDatasetInterface,\
    PascalAUG2012SegmentationDataSetInterface, PascalVOC2012SegmentationDataSetInterface
from super_gradients.training.datasets.dataset_interfaces.dataset_interface import \
    ClassificationTestDatasetInterface, CityscapesDatasetInterface, CoCoDetectionDatasetInterface


class DatasetsFactory(BaseFactory):

    def __init__(self):
        type_dict = {
            "classification_test_dataset": ClassificationTestDatasetInterface,
            "library_dataset": LibraryDatasetInterface,
            "classification_dataset": ClassificationDatasetInterface,
            "cifar_10": Cifar10DatasetInterface,
            "cifar_100": Cifar100DatasetInterface,
            "imagenet": ImageNetDatasetInterface,
            "tiny_imagenet": TinyImageNetDatasetInterface,
            "coco2017_detection": CoCoDetectionDatasetInterface,
            "coco2017_segmentation": CoCoSegmentationDatasetInterface,
            "pascal_voc_segmentation": PascalVOC2012SegmentationDataSetInterface,
            "pascal_aug_segmentation": PascalAUG2012SegmentationDataSetInterface,
            "cityscapes": CityscapesDatasetInterface,
        }
        super().__init__(type_dict)

@@ -49,9 +49,7 @@ phase_callbacks = [LRSchedulerCallback(scheduler=rop_lr_scheduler, phase=Phase.V
                    LRSchedulerCallback(scheduler=step_lr_scheduler, phase=Phase.TRAIN_EPOCH_END)]
 
 # Bring everything together with Trainer and start training
-trainer = Trainer("Cifar10_external_objects_example", multi_gpu=MultiGPUMode.OFF,
-                  train_loader=train_loader, valid_loader=valid_loader, classes=train_dataset.classes)
-trainer.build_model(net)
+trainer = Trainer("Cifar10_external_objects_example", multi_gpu=MultiGPUMode.OFF)
 
 train_params = {"max_epochs": 300,
                 "phase_callbacks": phase_callbacks,
@@ -65,4 +63,4 @@ train_params = {"max_epochs": 300,
                 "greater_metric_to_watch_is_better": True,
                 "lr_scheduler_step_type": "epoch"}
 
-trainer.train(training_params=train_params)
+trainer.train(model=net, training_params=train_params, train_loader=train_loader, valid_loader=valid_loader)

@@ -14,15 +14,14 @@ Paper:              https://arxiv.org/pdf/2101.06085.pdf
 
 import torch
 
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface
-
+from super_gradients.common import MultiGPUMode
+from super_gradients.training.datasets.datasets_utils import RandomResizedCropAndInterpolation
+from torchvision.transforms import RandomHorizontalFlip, ColorJitter, ToTensor, Normalize
 import super_gradients
-from super_gradients.training import Trainer, MultiGPUMode
-from super_gradients.training.models import HpmStruct
+from super_gradients.training import Trainer, models, dataloaders
 import argparse
-
 from super_gradients.training.metrics import Accuracy, Top5
-
+from super_gradients.training.datasets.data_augmentation import RandomErase
 parser = argparse.ArgumentParser()
 super_gradients.init_trainer()
 
@@ -56,19 +55,27 @@ dataset_params = {"batch_size": args.batch,
                   "random_erase_prob": 0.2,
                   "random_erase_value": 'random',
                   "train_interpolation": 'random',
-                  "auto_augment_config_string": 'rand-m9-mstd0.5'
                   }
 
+
+train_transforms = [RandomResizedCropAndInterpolation(size=224, interpolation="random"),
+                    RandomHorizontalFlip(),
+                    ColorJitter(0.4, 0.4, 0.4),
+                    ToTensor(),
+                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                    RandomErase(0.2, "random")
+                    ]
+
 trainer = Trainer(experiment_name=args.experiment_name,
                   multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL,
                   device='cuda')
 
-dataset = ImageNetDatasetInterface(dataset_params=dataset_params)
-
-trainer.connect_dataset_interface(dataset, data_loader_num_workers=8 * devices)
+train_loader = dataloaders.imagenet_train(dataset_params={"transforms": train_transforms},
+                                          dataloader_params={"batch_size": args.batch})
+valid_loader = dataloaders.imagenet_val()
 
-arch_params = HpmStruct(**{"num_classes": 1000, "aux_head": False, "classification_mode": True, 'dropout_prob': 0.3})
+model = models.get("ddrnet_23_slim" if args.slim else "ddrnet_23",
+                   arch_params={"aux_head": False, "classification_mode": True, 'dropout_prob': 0.3},
+                   num_classes=1000)
 
-trainer.build_model(architecture="ddrnet_23_slim" if args.slim else "ddrnet_23",
-                    arch_params=arch_params)
-trainer.train(training_params=train_params_ddr)
+trainer.train(model=model, training_params=train_params_ddr, train_loader=train_loader, valid_loader=valid_loader)

@@ -4,7 +4,8 @@ Deci-lab model export example.
 The main purpose of this code is to demonstrate how to upload the model to the platform, optimize and download it
  after training is complete, using DeciPlatformCallback.
 """
-from super_gradients import Trainer, ClassificationTestDatasetInterface
+from super_gradients import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.utils.callbacks import DeciLabUploadCallback, ModelConversionCheckCallback
 from deci_lab_client.models import (
@@ -28,8 +29,6 @@ def main(architecture_name: str):
         model_checkpoints_location="local",
         ckpt_root_dir=checkpoint_dir,
     )
-    dataset = ClassificationTestDatasetInterface(dataset_params={"batch_size": 10})
-    trainer.connect_dataset_interface(dataset, data_loader_num_workers=0)
 
     trainer.build_model(architecture=architecture_name, arch_params={"use_aux_heads": True, "aux_head": True})
 
@@ -91,7 +90,8 @@ def main(architecture_name: str):
 
     # RUN TRAINING. ONCE ALL EPOCHS ARE DONE THE OPTIMIZED MODEL FILE WILL BE LOCATED IN THE EXPERIMENT'S
     # CHECKPOINT DIRECTORY
-    trainer.train(train_params)
+    trainer.train(train_params, train_loader=classification_test_dataloader(),
+                  valid_loader=classification_test_dataloader())
 
 
 if __name__ == "__main__":

@@ -1,15 +1,12 @@
 import os
-from super_gradients.training import Trainer
-from super_gradients.training.datasets.dataset_interfaces import Cifar10DatasetInterface
+from super_gradients.training import Trainer, models
 from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
-
+from super_gradients.training.dataloaders.dataloaders import cifar10_train, cifar10_val
 os.environ["DECI_PLATFORM_TOKEN"] = "XXX"  # Replace XXX with your token
 
 
 trainer = Trainer(experiment_name='demo-deci-platform-logger')
-dataset = Cifar10DatasetInterface(dataset_params={"batch_size": 256, "val_batch_size": 512})
-trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
-trainer.build_model("resnet18")
+model = models.get("resnet18", num_classes=10)
 
 trainer.train(training_params={"max_epochs": 20,
                                "lr_updates": [5, 10, 15],
@@ -23,4 +20,6 @@ trainer.train(training_params={"max_epochs": 20,
                                "valid_metrics_list": [Accuracy(), Top5()],
                                "metric_to_watch": "Accuracy",
                                "greater_metric_to_watch_is_better": True,
-                               "sg_logger": "deci_platform_sg_logger"})
+                               "sg_logger": "deci_platform_sg_logger"},
+              train_loader=cifar10_train(),
+              valid_loader=cifar10_val())

@@ -2,10 +2,11 @@
 # Reaches ~94.9 Accuracy after 250 Epochs
 import super_gradients
 from super_gradients import Trainer
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import Cifar10DatasetInterface
+from super_gradients.training import models, dataloaders
 from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
 from super_gradients.training.utils.early_stopping import EarlyStop
 from super_gradients.training.utils.callbacks import Phase
+
 # Define Parameters
 super_gradients.init_trainer()
 
@@ -22,10 +23,8 @@ train_params = {"max_epochs": 250, "lr_updates": [100, 150, 200], "lr_decay_fact
 # Define Model
 trainer = Trainer("Callback_Example")
 
-# Connect Dataset
-dataset = Cifar10DatasetInterface()
-trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
-
 # Build Model
-trainer.build_model("resnet18_cifar")
-trainer.train(training_params=train_params)
+model = models.get("resnet18_cifar", num_classes=10)
+
+trainer.train(model=model, training_params=train_params,
+              train_loader=dataloaders.cifar10_train(), valid_loader=dataloaders.cifar10_val())

# Cifar10 Classification Training:
# Reaches ~94.9 Accuracy after 250 Epochs

import super_gradients
from omegaconf import DictConfig
import hydra
import pkg_resources


@hydra.main(config_path=pkg_resources.resource_filename("conf", ""), config_name="cifar10_resnet_conf")
def train(cfg: DictConfig) -> None:
    # INSTANTIATE ALL OBJECTS IN CFG
    cfg = hydra.utils.instantiate(cfg)

    # CONNECT THE DATASET INTERFACE WITH DECI MODEL
    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

    # BUILD NETWORK
    cfg.trainer .build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)

    # TRAIN
    cfg.trainer.train(training_params=cfg.training_params)


if __name__ == "__main__":
    super_gradients.init_trainer()
    train()

# Darknet53 Backbone Training on HAM10000 Dataset
from super_gradients.training import MultiGPUMode
from super_gradients.training import Trainer
from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationDatasetInterface

# Define Parameters
train_params = {"max_epochs": 110, "lr_updates": [30, 60, 90, 100], "lr_decay_factor": 0.1, "lr_mode": "step",
                "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD",
                "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}}
arch_params = {'backbone_mode': False, 'num_classes': 7}
dataset_params = {"batch_size": 16, "test_batch_size": 16, 'dataset_dir': '/data/HAM10000'}

# Define Model
trainer = Trainer("Darknet53_Backbone_HAM10000",
                  model_checkpoints_location='local',
                  device='cuda',
                  multi_gpu=MultiGPUMode.DATA_PARALLEL)

# Connect Dataset
dataset = ClassificationDatasetInterface(normalization_mean=(0.7483, 0.5154, 0.5353),
                                         normalization_std=(0.1455, 0.1691, 0.1879),
                                         resolution=416,
                                         dataset_params=dataset_params)

trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)

# Build Model
trainer.build_model("darknet53", arch_params=arch_params)

# Start Training
trainer.train(training_params=train_params)

"""EfficientNet-b0 training on Imagenet
TODO: This example code is the STARTING POINT for training EfficientNet - IT DIDN'T ACHIEVE THE PAPER'S ACCURACY!!!
Training params are set according to https://github.com/rwightman/pytorch-image-models/issues/11
Training on 4 GPUs with initial LR = 0.0032 achieves ~74.7%, (Paper=77.1% Timm=77.69%)
The Tensorboards of the previous attempts: 's3/deci-model-repository-research/enet_reproduce_attempts'
"""

import super_gradients
from omegaconf import DictConfig
import hydra
import pkg_resources


@hydra.main(config_path=pkg_resources.resource_filename("conf", ""), config_name="imagenet_efficientnet_conf")
def train(cfg: DictConfig) -> None:
    # INSTANTIATE ALL OBJECTS IN CFG
    cfg = hydra.utils.instantiate(cfg)

    # CONNECT THE DATASET INTERFACE WITH DECI MODEL
    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

    # BUILD NETWORK
    cfg.trainer .build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)

    # TRAIN
    cfg.trainer.train(training_params=cfg.training_params)


if __name__ == "__main__":
    super_gradients.init_trainer()
    train()

# MobileNetV3 Large Imagenet classification training:
# This example trains with batch_size = 128 * 2 GPUs, total 256.
# Training time on 2 X GeForce RTX 2080 Ti is 19min / epoch, total time ~ 50 hours.
# Reach 73.79 Top1 accuracy.
# Training parameters are for MobileNet Large

import super_gradients
from omegaconf import DictConfig
import hydra
import pkg_resources


@hydra.main(config_path=pkg_resources.resource_filename("conf", ""), config_name="imagenet_mobilenetv3_conf")
def train(cfg: DictConfig) -> None:
    # INSTANTIATE ALL OBJECTS IN CFG
    cfg = hydra.utils.instantiate(cfg)

    # CONNECT THE DATASET INTERFACE WITH DECI MODEL
    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

    # BUILD NETWORK
    cfg.trainer .build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)

    # TRAIN
    cfg.trainer.train(training_params=cfg.training_params)


if __name__ == "__main__":
    super_gradients.init_trainer()
    train()

# Imagenet classification training:
# For RegnetY800 => 76.1 accuracy

# the hyper-parameters are tailored for training on Single 2080Ti GPU.

import super_gradients
from omegaconf import DictConfig
import hydra
import pkg_resources


@hydra.main(config_path=pkg_resources.resource_filename("conf", ""), config_name="imagenet_regnetY800_conf")
def train(cfg: DictConfig) -> None:
    # INSTANTIATE ALL OBJECTS IN CFG
    cfg = hydra.utils.instantiate(cfg)

    # CONNECT THE DATASET INTERFACE WITH DECI MODEL
    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

    # BUILD NETWORK
    cfg.trainer .build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)

    # TRAIN
    cfg.trainer.train(training_params=cfg.training_params)


if __name__ == "__main__":
    super_gradients.init_trainer()
    train()

import super_gradients
from omegaconf import DictConfig
import hydra
import pkg_resources


@hydra.main(config_path=pkg_resources.resource_filename("conf", ""), config_name="imagenet_repvgg_conf")
def train(cfg: DictConfig) -> None:
    # INSTANTIATE ALL OBJECTS IN CFG
    cfg = hydra.utils.instantiate(cfg)

    # CONNECT THE DATASET INTERFACE WITH DECI MODEL
    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

    # BUILD NETWORK
    cfg.trainer .build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)

    # TRAIN
    cfg.trainer.train(training_params=cfg.training_params)


if __name__ == "__main__":
    super_gradients.init_trainer()
    train()

"""
ResNet50 Imagenet classification training:
This example trains with batch_size = 64 * 4 GPUs, total 256.

Training times:
    ResNet18:   36 hours with 4 X NVIDIA RTX A5000.
    ResNet34:   36 hours with 4 X NVIDIA RTX A5000.
    ResNet50:   46 hours with 4 X GeForce RTX 3090 Ti.

Top1, Top5 results:
    ResNet18:   Top1: 70.60  Top5: 89.64
    ResNet34:   Top1: 74.13  Top5: 91.70
    ResNet50:   Top1: 76.30  Top5: 93.03

BE AWARE THAT THIS RECIPE USE DATA_PARALLEL, WHEN USING DDP FOR DISTRIBUTED TRAINING THIS RECIPE REACH ONLY 75.4 TOP1
ACCURACY.
"""

import super_gradients
from omegaconf import DictConfig
import hydra
import pkg_resources


@hydra.main(config_path=pkg_resources.resource_filename("conf", ""), config_name="imagenet_resnet50_conf")
def train(cfg: DictConfig) -> None:
    # INSTANTIATE ALL OBJECTS IN CFG
    cfg = hydra.utils.instantiate(cfg)

    # CONNECT THE DATASET INTERFACE WITH DECI MODEL
    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

    # BUILD NETWORK
    cfg.trainer .build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)

    # TRAIN
    cfg.trainer.train(training_params=cfg.training_params)


if __name__ == "__main__":
    super_gradients.init_trainer()
    train()

#!/usr/bin/env python
""" Single node distributed training.

    The program will dispatch distributed training on all available GPUs residing in a single node.

    Usage:
    python -m torch.distributed.launch --nproc_per_node=n distributed_training_imagenet.py
    where n is the number of GPUs required, e.g., n=8

    Important note: (1) in distributed training it is customary to specify learning rates and batch sizes per GPU.
    Whatever learning rate and schedule you specify will be applied to the each GPU individually.
    Since gradients are passed and summed (reduced) from all to all GPUs, the effective batch size is the
    batch you specify times the number of GPUs. In the literature there are several "best practices" to set
    learning rates and schedules for large batch sizes.
    Should be checked with. (2) The training protocol specified in this file for 8 GPUs are far from optimal.
    The best protocol should use cosine schedule.

    In the example below: for ImageNet training using Resnet50, when applied with n=8 should compute an Eopch in about
    5min20sec with 8 V100 GPUs.

    Todo: (1) the code is more or less ready for multiple nodes, but I have not experimented with it at all.
          (2) detection and segmentation codes were not modified and should not work properly.
              Specifically, the analogue changes done in sg_classification_model should be done also in
              deci_segmentation_model and deci_detection_model

"""
import super_gradients
import torch.distributed
from super_gradients.training.sg_trainer import MultiGPUMode
from super_gradients.training import Trainer
from super_gradients.training.datasets.dataset_interfaces import ImageNetDatasetInterface
from super_gradients.common.aws_connection.aws_secrets_manager_connector import AWSSecretsManagerConnector
from super_gradients.training.metrics.classification_metrics import Accuracy, Top5

torch.backends.cudnn.benchmark = True

super_gradients.init_trainer()
# TODO - VALIDATE THE HYPER PARAMETERS WITH RAN TO FIX THIS EXAMPLE CODE
train_params = {"max_epochs": 110,
                "lr_updates": [30, 60, 90],
                "lr_decay_factor": 0.1,
                "initial_lr": 0.6,
                "loss": "cross_entropy",
                "lr_mode": "step",
                # "initial_lr": 0.05 * 2,
                "lr_warmup_epochs": 5,
                # "criterion_params":{"smooth_eps":0.1}}
                "mixed_precision": True,
                # "mixed_precision_opt_level": "O3",
                "optimizer_params": {"weight_decay": 0.000, "momentum": 0.9},
                # "optimizer_params": {"weight_decay": 0.0001, "momentum": 0.9}
                "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()],
                "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                "greater_metric_to_watch_is_better": True}
dataset_params = {"batch_size": 128}

model_repo_bucket_name = AWSSecretsManagerConnector.get_secret_value_for_secret_key(aws_env='research',
                                                                                    secret_name='training_secrets',
                                                                                    secret_key='S3.MODEL_REPOSITORY_BUCKET_NAME')
trainer = Trainer("test_checkpoints_resnet_8_gpus",
                  model_checkpoints_location='s3://' + model_repo_bucket_name,
                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL
                  )
# FOR AWS
dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params=dataset_params)
trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
trainer.build_model("resnet50")
trainer.train(training_params=train_params)

# ShelfNet LW 34 training on CoCo Segmentation Dataset:
# mIOU on CoCo Seg: ~0.66

# Since the code is training on a Subset of COCO Seg, there is an initial creation process for the "Sub-DataSet"
# this training process is optimized to enable fine-tuning on PASCAL VOC 2012 Dataset that has only 21 Classes...

# IMPORTANT: The code is optimized for a fixed initial LR since the Polynomial Loss is pretty sensitive, so we keep the
# same LR by dividing by the number of GPUs (since our code base multiplies it automatically)

# P.S. - Use the relevant training params dict if you are running on TZAG or on V100

import torch
from super_gradients.training import Trainer, MultiGPUMode
from super_gradients.training.datasets import CoCoSegmentationDatasetInterface
from super_gradients.training.sg_trainer.sg_trainer import StrictLoad
from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU

model_size_str = '34'

coco_sub_classes_inclusion_tuples_list = [(0, 'background'), (5, 'airplane'), (2, 'bicycle'), (16, 'bird'),
                                          (9, 'boat'),
                                          (44, 'bottle'), (6, 'bus'), (3, 'car'), (17, 'cat'), (62, 'chair'),
                                          (21, 'cow'),
                                          (67, 'dining table'), (18, 'dog'), (19, 'horse'), (4, 'motorcycle'),
                                          (1, 'person'),
                                          (64, 'potted plant'), (20, 'sheep'), (63, 'couch'), (7, 'train'),
                                          (72, 'tv')]

coco_seg_dataset_tzag_params = {
    "batch_size": 24,
    "test_batch_size": 24,
    "dataset_dir": "/data/coco/",
    "s3_link": None,
    "img_size": 608,
    "crop_size": 512
}

coco_seg_dataset_v100_params = {
    "batch_size": 32,
    "test_batch_size": 32,
    "dataset_dir": "/home/ubuntu/data/coco/",
    "s3_link": None,
    "img_size": 608,
    "crop_size": 512
}

shelfnet_coco_training_params = {
    "max_epochs": 150, "initial_lr": 5e-3, "loss": "shelfnet_ohem_loss",
    "optimizer": "SGD", "mixed_precision": True, "lr_mode": "poly",
    "optimizer_params": {"momentum": 0.9, "weight_decay": 1e-4, "nesterov": False},
    "load_opt_params": False, "train_metrics_list": [PixelAccuracy(), IoU(21)],
    "valid_metrics_list": [PixelAccuracy(), IoU(21)],
    "loss_logging_items_names": ["Loss1/4", "Loss1/8", "Loss1/16", "Loss"], "metric_to_watch": "IoU",
    "greater_metric_to_watch_is_better": True}

shelfnet_lw_arch_params = {"num_classes": 21, "load_checkpoint": True, "strict_load": StrictLoad.ON,
                           "multi_gpu_mode": "data_parallel", "load_weights_only": True,
                           "load_backbone": True, "source_ckpt_folder_name": 'resnet' + model_size_str}

data_loader_num_workers = 8 * torch.cuda.device_count()

# BUILD THE LIGHT-WEIGHT SHELFNET ARCHITECTURE FOR TRAINING
experiment_name_prefix = 'shelfnet_lw_'
experiment_name_dataset_suffix = '_coco_seg_' + str(
    shelfnet_coco_training_params['max_epochs']) + '_epochs_train_example'

experiment_name = experiment_name_prefix + model_size_str + experiment_name_dataset_suffix

trainer = Trainer(experiment_name,
                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL,
                  ckpt_name='ckpt_best.pth')

coco_seg_datasaet_interface = CoCoSegmentationDatasetInterface(dataset_params=coco_seg_dataset_tzag_params,
                                                               cache_labels=False,
                                                               dataset_classes_inclusion_tuples_list=coco_sub_classes_inclusion_tuples_list)

trainer.connect_dataset_interface(coco_seg_datasaet_interface, data_loader_num_workers=data_loader_num_workers)
trainer.build_model('shelfnet' + model_size_str, arch_params=shelfnet_lw_arch_params)

print('Training ShelfNet-LW model: ' + experiment_name)
trainer.train(training_params=shelfnet_coco_training_params)

@@ -1,32 +1,29 @@
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import SuperviselyPersonsDatasetInterface
+from super_gradients.training import models, dataloaders
+
 from super_gradients.training.sg_trainer import Trainer
 from super_gradients.training.metrics import BinaryIOU
 from super_gradients.training.transforms.transforms import ResizeSeg, RandomFlip, RandomRescale, CropImageAndMask, \
     PadShortToCropSize, ColorJitterSeg
 from super_gradients.training.utils.callbacks import BinarySegmentationVisualizationCallback, Phase
-from torchvision import transforms
 
 # DEFINE DATA TRANSFORMATIONS
-dataset_params = {
-    "image_mask_transforms_aug": transforms.Compose([ColorJitterSeg(brightness=0.5, contrast=0.5, saturation=0.5),
-                                                     RandomFlip(),
-                                                     RandomRescale(scales=[0.25, 1.]),
-                                                     PadShortToCropSize([320, 480]),
-                                                     CropImageAndMask(crop_size=[320, 480],
-                                                                      mode="random")]),
-    "image_mask_transforms": transforms.Compose([ResizeSeg(h=480, w=320)])
-}
-
-dataset_interface = SuperviselyPersonsDatasetInterface(dataset_params)
 
-trainer = Trainer("regseg48_transfer_learning_old_dice_diff_lrs_head_fixed_50_epochs")
+dl_train = dataloaders.supervisely_persons_train(
+    dataset_params={"transforms": [ColorJitterSeg(brightness=0.5, contrast=0.5, saturation=0.5),
+                                   RandomFlip(),
+                                   RandomRescale(scales=[0.25, 1.]),
+                                   PadShortToCropSize([320, 480]),
+                                   CropImageAndMask(crop_size=[320, 480],
+                                                    mode="random")]})
 
-# CONNECTING THE DATASET INTERFACE WILL SET SGMODEL'S CLASSES ATTRIBUTE ACCORDING TO SUPERVISELY
-trainer.connect_dataset_interface(dataset_interface)
+dl_val = dataloaders.supervisely_persons_val(dataset_params={"transforms": [ResizeSeg(h=480, w=320)]})
+
+trainer = Trainer("regseg48_transfer_learning_old_dice_diff_lrs_head_fixed_50_epochs")
 
 # THIS IS WHERE THE MAGIC HAPPENS- SINCE SGMODEL'S CLASSES ATTRIBUTE WAS SET TO BE DIFFERENT FROM CITYSCAPES'S, AFTER
 # LOADING THE PRETRAINED REGSET, IT WILL CALL IT'S REPLACE_HEAD METHOD AND CHANGE IT'S SEGMENTATION HEAD LAYER ACCORDING
 # TO OUR BINARY SEGMENTATION DATASET
+model = models.get("regseg48", pretrained_weights="cityscapes", num_classes=1)
 trainer.build_model("regseg48", arch_params={"pretrained_weights": "cityscapes"})
 
 # DEFINE TRAINING PARAMS. SEE DOCS FOR THE FULL LIST.

@@ -16,9 +16,7 @@ Once triggered, the following will happen:
 Finally, once training is over- we trigger a pos-training callback that will export the ONNX files.
 
 """
-
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface
-from super_gradients.training import Trainer, MultiGPUMode
+from super_gradients.training import Trainer, MultiGPUMode, models, dataloaders
 from super_gradients.training.metrics.classification_metrics import Accuracy
 
 import super_gradients
@@ -26,13 +24,14 @@ from super_gradients.training.utils.quantization_utils import PostQATConversionC
 
 super_gradients.init_trainer()
 
-dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params={"batch_size": 128})
 trainer = Trainer("resnet18_qat_example",
                   model_checkpoints_location='local',
                   multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)
 
-trainer.connect_dataset_interface(dataset)
-trainer.build_model("resnet18", checkpoint_params={"pretrained_weights": "imagenet"})
+train_loader = dataloaders.imagenet_train()
+valid_loader = dataloaders.imagenet_val()
+
+model = models.get("resnet18", pretrained_weights="imagenet")
 
 train_params = {"max_epochs": 1,
                 "lr_mode": "step",
@@ -58,4 +57,4 @@ train_params = {"max_epochs": 1,
                 "phase_callbacks": [PostQATConversionCallback(dummy_input_size=(1, 3, 224, 224))]
                 }
 
-trainer.train(training_params=train_params)
+trainer.train(model=model, training_params=train_params, train_loader=train_loader, valid_loader=valid_loader)

# TODO: REFACTOR AS YAML FILES RECIPE
import super_gradients
import torch
from super_gradients.training.datasets import PascalAUG2012SegmentationDataSetInterface
from super_gradients.training import Trainer, MultiGPUMode
from super_gradients.training.sg_trainer.sg_trainer import StrictLoad
from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU

super_gradients.init_trainer()
pascal_aug_dataset_params = {"batch_size": 16,
                             "test_batch_size": 16,
                             "dataset_dir": "/data/pascal_voc_2012/VOCaug/dataset/",
                             "s3_link": None,
                             "img_size": 512,
                             "train_loader_drop_last": True,
                             }
shelfnet_lw_pascal_aug_training_params = {"max_epochs": 250, "initial_lr": 1e-2, "loss": "shelfnet_ohem_loss",
                                          "optimizer": "SGD", "mixed_precision": False, "lr_mode": "poly",
                                          "optimizer_params": {"momentum": 0.9, "weight_decay": 1e-4,
                                                               "nesterov": False},
                                          "load_opt_params": False, "train_metrics_list": [PixelAccuracy(), IoU(21)],
                                          "valid_metrics_list": [PixelAccuracy(), IoU(21)],
                                          "loss_logging_items_names": ["Loss1/4", "Loss1/8", "Loss1/16", "Loss"],
                                          "metric_to_watch": "IoU",
                                          "greater_metric_to_watch_is_better": True}

shelfnet_lw_arch_params = {"num_classes": 21, "strict_load": StrictLoad.ON,
                           "multi_gpu_mode": MultiGPUMode.OFF}

checkpoint_params = {"load_checkpoint": True, "load_weights_only": True,
                     "load_backbone": True, "source_ckpt_folder_name": 'resnet_backbones'}

if torch.cuda.is_available() and torch.cuda.device_count() > 1:
    data_loader_num_workers = 16
    shelfnet_lw_pascal_aug_training_params["initial_lr"] = shelfnet_lw_pascal_aug_training_params["initial_lr"] / 2.
else:
    # SINGLE GPU TRAINING
    data_loader_num_workers = 8

# SET THE *LIGHT-WEIGHT* SHELFNET ARCHITECTURE SIZE (UN-COMMENT TO TRAIN)
model_size_str = '34'
# model_size_str = '18'
# BUILD THE LIGHT-WEIGHT SHELFNET ARCHITECTURE FOR TRAINING
experiment_name_prefix = 'shelfnet_lw_'
experiment_name_dataset_suffix = '_pascal_aug_encoding_dataset_train_250_epochs_no_batchnorm_decoder'
experiment_name = experiment_name_prefix + model_size_str + experiment_name_dataset_suffix
trainer = Trainer(experiment_name, model_checkpoints_location='local', multi_gpu=True,
                  ckpt_name='resnet' + model_size_str + '.pth')

pascal_aug_datasaet_interface = PascalAUG2012SegmentationDataSetInterface(
    dataset_params=pascal_aug_dataset_params,
    cache_labels=False)
trainer.connect_dataset_interface(pascal_aug_datasaet_interface, data_loader_num_workers=data_loader_num_workers)
trainer.build_model('shelfnet' + model_size_str, arch_params=shelfnet_lw_arch_params, checkpoint_params=checkpoint_params)
print('Training ShelfNet-LW model: ' + experiment_name)
trainer.train(training_params=shelfnet_lw_pascal_aug_training_params)

"""
This file is used to define the Dataset used for the Training.
"""
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from super_gradients.training import utils as core_utils
from super_gradients.training.datasets.dataset_interfaces import DatasetInterface


class UserDataset(DatasetInterface):
    """
    The user's dataset inherits from SuperGradient's DatasetInterface and must
    contain a trainset and testset from which the the data will be loaded using.
    All augmentations, resizing and parsing must be done in this class.

     - Augmentations are defined below and will be carried out in the order they are given.
       super_gradients provides additional dataset reading tools such as ListDataset given a list of files
       corresponding to the images and labels.

    """
    def __init__(self, name="cifar10", dataset_params={}):
        super(UserDataset, self).__init__(dataset_params)
        self.dataset_name = name
        self.lib_dataset_params = {'mean': (0.4914, 0.4822, 0.4465), 'std': (0.2023, 0.1994, 0.2010)}

        crop_size = core_utils.get_param(self.dataset_params, 'crop_size', default_val=32)

        transform_train = transforms.Compose([
            transforms.RandomCrop(crop_size, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(self.lib_dataset_params['mean'], self.lib_dataset_params['std']),
        ])

        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(self.lib_dataset_params['mean'], self.lib_dataset_params['std']),
        ])

        self.trainset = datasets.CIFAR10(root=self.dataset_params.dataset_dir, train=True, download=True,
                                         transform=transform_train)

        self.testset = datasets.CIFAR10(root=self.dataset_params.dataset_dir, train=False, download=True,
                                        transform=transform_test)

"""
The loss must be of torch.nn.modules.loss._Loss class.
For commonly used losses, import from deci.core.ADNN.losses

-IMPORTANT: forward(...) should return (loss, loss_items) where loss is the tensor used for backprop (i.e what your
original loss function returns), and loss_items should be a tensor of shape (n_items), of values computed during
the forward pass which we desire to log over the entire epoch. For example- the loss itself should always be logged.
Another examploe is a scenario where the computed loss is the sum of a few components we would like to log- these
entries in loss_items).

-When training, set the "loss_logging_items_names" parameter in train_params to be a list of strings, of length
n_items who's ith element is the name of the ith entry in loss_items. Then each item will be logged, rendered on
tensorboard and "watched" (i.e saving model checkpoints according to it).

-Since running logs will save the loss_items in some internal state, it is recommended that loss_items are detached
from their computational graph for memory efficiency.
"""

import torch.nn as nn
from super_gradients.training.losses.label_smoothing_cross_entropy_loss import cross_entropy


class LabelSmoothingCrossEntropyLoss(nn.CrossEntropyLoss):
    """
    LabelSmoothingCrossEntropyLoss - POC loss class, uses SuperGradient's cross entropy which support distribution as targets.

    """

    def __init__(self, weight=None, ignore_index=-100, reduction='mean', smooth_eps=None, smooth_dist=None,
                 from_logits=True):
        super(LabelSmoothingCrossEntropyLoss, self).__init__(weight=weight,
                                                             ignore_index=ignore_index, reduction=reduction)
        self.smooth_eps = smooth_eps
        self.smooth_dist = smooth_dist
        self.from_logits = from_logits

    def forward(self, input, target, smooth_dist=None):
        if smooth_dist is None:
            smooth_dist = self.smooth_dist
        loss = cross_entropy(input, target, weight=self.weight, ignore_index=self.ignore_index,
                             reduction=self.reduction, smooth_eps=self.smooth_eps,
                             smooth_dist=smooth_dist, from_logits=self.from_logits)

        loss_items = loss.detach().unsqueeze(0)

        return loss, loss_items

"""
This file is used to define the Metrics used for training.
The metrics object must be of torchmetrics.Metric type. For more information on how to use torchmetric.Metric objects and
 implement your own metrics see https://torchmetrics.readthedocs.io/en/latest/pages/overview.html
"""

import torchmetrics
import torch


class Accuracy(torchmetrics.Accuracy):
    def __init__(self, dist_sync_on_step=False):
        super().__init__(dist_sync_on_step=dist_sync_on_step, top_k=1)

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        super().update(preds=preds.softmax(1), target=target)


class Top5(torchmetrics.Accuracy):
    def __init__(self, dist_sync_on_step=False):
        super().__init__(dist_sync_on_step=dist_sync_on_step, top_k=5)

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        super().update(preds=preds.softmax(1), target=target)

"""
This file is used to define the model used for training. For example, in this template, we define ResNet50.
One may use existing models from torchvision as well (e.g., torchvision.models.resnet50)
"""

import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion * planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_planes, planes, stride=1):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.expansion * planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion * planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


def width_multiplier(original, factor):
    return int(original * factor)


class ResNet(nn.Module):
    def __init__(self, block, num_blocks: list, num_classes: int = 10, width_mult: float = 1,
                 input_batchnorm: bool = False, backbone_mode: bool = False):
        super(ResNet, self).__init__()
        self.backbone_mode = backbone_mode
        self.structure = [num_blocks, width_mult]
        self.in_planes = width_multiplier(64, width_mult)
        self.input_batchnorm = input_batchnorm
        if self.input_batchnorm:
            self.bn0 = nn.BatchNorm2d(3)

        self.conv1 = nn.Conv2d(3, width_multiplier(64, width_mult), kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(width_multiplier(64, width_mult))
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(block, width_multiplier(64, width_mult), num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, width_multiplier(128, width_mult), num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, width_multiplier(256, width_mult), num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, width_multiplier(512, width_mult), num_blocks[3], stride=2)

        if not self.backbone_mode:
            # IF RESNET IS IN BACK_BONE MODE WE DON'T NEED THE FINAL CLASSIFIER LAYERS, BUT ONLY THE NET BLOCK STRUCTURE
            self.linear = nn.Linear(width_multiplier(512, width_mult) * block.expansion, num_classes)
            self.avgpool = nn.AdaptiveAvgPool2d(1)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        if num_blocks == 0:
            # When the number of blocks is zero but spatial dimension and/or number of filters about to change we put 1
            # 3X3 conv layer to make this change to the new dimensions.
            if stride != 1 or self.in_planes != planes:
                layers.append(nn.Sequential(
                    nn.Conv2d(self.in_planes, planes, kernel_size=3, stride=stride, bias=False, padding=1),
                    nn.BatchNorm2d(planes))
                )
                self.in_planes = planes

        else:
            for stride in strides:
                layers.append(block(self.in_planes, planes, stride))
                self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        if self.input_batchnorm:
            x = self.bn0(x)
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.maxpool(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)

        if not self.backbone_mode:
            # IF RESNET IS *NOT* IN BACK_BONE MODE WE  NEED THE FINAL CLASSIFIER LAYERS OUTPUTS
            out = self.avgpool(out)
            out = out.squeeze(dim=2).squeeze(dim=2)
            out = self.linear(out)

        return out

    def load_state_dict(self, state_dict, strict=True):
        """
        load_state_dict - Overloads the base method and calls it to load a modified dict for usage as a backbone
        :param state_dict:  The state_dict to load
        :param strict:      strict loading (see super() docs)
        """
        pretrained_model_weights_dict = state_dict.copy()

        if self.backbone_mode:
            # FIRST LET'S POP THE LAST TWO LAYERS - NO NEED TO LOAD THEIR VALUES SINCE THEY ARE IRRELEVANT AS A BACKBONE
            pretrained_model_weights_dict.popitem()
            pretrained_model_weights_dict.popitem()

            pretrained_backbone_weights_dict = OrderedDict()
            for layer_name, weights in pretrained_model_weights_dict.items():
                # GET THE LAYER NAME WITHOUT THE 'module.' PREFIX
                name_without_module_prefix = layer_name.split('module.')[1]

                # MAKE SURE THESE ARE NOT THE FINAL LAYERS
                pretrained_backbone_weights_dict[name_without_module_prefix] = weights

            # RETURNING THE UNMODIFIED/MODIFIED STATE DICT DEPENDING ON THE backbone_mode VALUE
            super().load_state_dict(pretrained_backbone_weights_dict, strict)
        else:
            super().load_state_dict(pretrained_model_weights_dict, strict)

from super_gradients.training import Trainer
from super_gradients.training import MultiGPUMode
from dataset import UserDataset
from model import ResNet, BasicBlock
from loss import LabelSmoothingCrossEntropyLoss
from metrics import Accuracy, Top5


def main():
    # ------------------ Loading The Model From Model.py----------------
    arch_params = {'num_classes': 10}
    model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=arch_params['num_classes'])

    trainer = Trainer('client_model_training',
                      model_checkpoints_location='local',
                      multi_gpu=MultiGPUMode.OFF)

    # if a torch.nn.Module is provided when building the model, the model will be integrated into deci model class
    trainer.build_model(model, arch_params=arch_params)

    # ------------------ Loading The Dataset From Dataset.py----------------
    dataset_params = {"batch_size": 256}
    dataset = UserDataset(dataset_params)
    trainer.connect_dataset_interface(dataset)

    # ------------------ Loading The Loss From Loss.py -----------------
    loss = LabelSmoothingCrossEntropyLoss()

    # ------------------ Defining the metrics we wish to log -----------------
    train_metrics_list = [Accuracy(), Top5()]
    valid_metrics_list = [Accuracy(), Top5()]

    # ------------------ Training -----------------
    train_params = {"max_epochs": 250,
                    "lr_updates": [100, 150, 200],
                    "lr_decay_factor": 0.1,
                    "lr_mode": "step",
                    "lr_warmup_epochs": 0,
                    "initial_lr": 0.1,
                    "loss": loss,
                    "criterion_params": {},
                    "optimizer": "SGD",
                    "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
                    "launch_tensorboard": False,
                    "train_metrics_list": train_metrics_list,
                    "valid_metrics_list": valid_metrics_list,
                    "loss_logging_items_names": ["Loss"],
                    "metric_to_watch": "Accuracy",
                    "greater_metric_to_watch_is_better": True}

    trainer.train(train_params)


if __name__ == '__main__':
    main()

@@ -12,9 +12,8 @@ defaults:
   - arch_params: resnet18_cifar_arch_params
   - checkpoint_params: default_checkpoint_params
 
-dataset_interface:
-  cifar_10:
-    dataset_params: ${dataset_params}
+train_dataloader: cifar10_train
+val_dataloader: cifar10_val
 
 data_loader_num_workers: 8
 
@@ -27,3 +26,5 @@ model_checkpoints_location: local
 ckpt_root_dir:
 
 architecture: resnet18_cifar
+
+experiment_name: resnet18_cifar

@@ -38,57 +38,12 @@
 
 defaults:
   - training_hyperparams: cityscapes_default_train_params
-  #  - dataset_params: cityscapes_ddrnet_dataset_params # TODO: uncomment after DatasetInterface refactor
+  - dataset_params: cityscapes_ddrnet_dataset_params
   - checkpoint_params: default_checkpoint_params
   - _self_
 
-# ------------------------------------- legacy dataset params ------------------------------------- #
-dataset_params:
-  _convert_: all
-  batch_size: 6
-  val_batch_size: 6
-  dataset_dir: /data/cityscapes
-  crop_size: [ 1024, 1024 ]
-  img_size: 1024
-  train_loader_drop_last: True
-  color_jitter: 0.5
-  random_scales: [ 0.5, 2. ]
-  eval_scale: 1.
-  cityscapes_ignored_label: 19
-
-  image_mask_transforms_aug:
-    Compose:
-      transforms:
-        - ColorJitterSeg:
-            brightness: ${dataset_params.color_jitter}
-            contrast: ${dataset_params.color_jitter}
-            saturation: ${dataset_params.color_jitter}
-
-        - RandomFlipSeg
-
-        - RandomRescaleSeg:
-            scales: ${dataset_params.random_scales}
-
-        - PadShortToCropSizeSeg:
-            crop_size: ${dataset_params.crop_size}
-            fill_mask: ${dataset_params.cityscapes_ignored_label}
-
-        - CropImageAndMaskSeg:
-            crop_size: ${dataset_params.crop_size}
-            mode: random
-
-  image_mask_transforms:
-    Compose:
-      transforms:
-        - RescaleSeg:
-            scale_factor: ${dataset_params.eval_scale}
-
-dataset_interface:
-  cityscapes:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 8
-# ------------------------------------- legacy dataset params ------------------------------------- #
+train_dataloader: cityscapes_train
+val_dataloader: cityscapes_val
 
 architecture: ddrnet_23
 
@@ -107,13 +62,13 @@ training_hyperparams:
       edge_kernel: 5
   loss_logging_items_names: [ main_loss, aux_loss1, loss ]
 
-
 arch_params:
   num_classes: 19
   aux_head: True
   sync_bn: True
 
 
+
 load_checkpoint: False
 checkpoint_params:
   load_checkpoint: ${load_checkpoint}
@@ -121,7 +76,6 @@ checkpoint_params:
   load_backbone: True
   strict_load: no_key_matching
 
-
 experiment_name: ${architecture}_cityscapes
 
 model_checkpoints_location: local

@@ -27,58 +27,12 @@
 
 defaults:
   - training_hyperparams: default_train_params
-  #  - dataset_params: cityscapes_regseg48_dataset_params # TODO: uncomment after DatasetInterface refactor
+  - dataset_params: cityscapes_regseg48_dataset_params
   - checkpoint_params: default_checkpoint_params
   - _self_
 
-# ------------------------------------- legacy dataset params ------------------------------------- #
-dataset_params:
-  _convert_: all
-  batch_size: 4
-  val_batch_size: 4
-  dataset_dir: /data/cityscapes
-  crop_size: 1024
-  img_size: 1024
-  train_loader_drop_last: True
-  color_jitter: 0.1
-  random_scales: [ 0.4, 1.6 ]
-  cityscapes_ignored_label: 19
-
-  image_mask_transforms_aug:
-    Compose:
-      transforms:
-        - ColorJitterSeg:
-            brightness: ${dataset_params.color_jitter}
-            contrast: ${dataset_params.color_jitter}
-            saturation: ${dataset_params.color_jitter}
-
-        - RandomFlipSeg
-
-        - RandomRescaleSeg:
-            scales: ${dataset_params.random_scales}
-
-        - PadShortToCropSizeSeg:
-            crop_size: ${dataset_params.crop_size}
-            fill_image:
-              - ${dataset_params.cityscapes_ignored_label}
-              - 0
-              - 0
-            fill_mask: ${dataset_params.cityscapes_ignored_label}
-
-        - CropImageAndMaskSeg:
-            crop_size: ${dataset_params.crop_size}
-            mode: random
-
-  image_mask_transforms:
-    Compose:
-      transforms: [ ]
-
-dataset_interface:
-  cityscapes:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 8
-# ------------------------------------- legacy dataset params ------------------------------------- #
+train_dataloader: cityscapes_train
+val_dataloader: cityscapes_val
 
 cityscapes_ignored_label: 19    # convenience parameter since it is used in many places in the YAML
 

@@ -5,39 +5,8 @@ defaults:
   - dataset_params: cityscapes_dataset_params
   - checkpoint_params: default_checkpoint_params
 
-
-dataset_params:
-  color_jitter: 0.5
-  image_mask_transforms_aug:
-    Compose:
-      transforms:
-        - ColorJitterSeg:
-            brightness: ${dataset_params.color_jitter}
-            contrast: ${dataset_params.color_jitter}
-            saturation: ${dataset_params.color_jitter}
-
-        - RandomFlipSeg
-
-        - RandomRescaleSeg:
-            scales: ${dataset_params.random_scales}
-
-        - PadShortToCropSizeSeg:
-            crop_size: ${dataset_params.crop_size}
-            fill_mask: ${dataset_params.cityscapes_ignored_label}
-
-        - CropImageAndMaskSeg:
-            crop_size: ${dataset_params.crop_size}
-            mode: random
-
-  image_mask_transforms:
-    Compose:
-      transforms:
-        - RescaleSeg:
-            scale_factor: ${dataset_params.eval_scale}
-
-dataset_interface:
-  cityscapes:
-    dataset_params: ${dataset_params}
+train_dataloader: cityscapes_train
+val_dataloader: cityscapes_val
 
 data_loader_num_workers: 10
 

@@ -39,62 +39,15 @@
 
 defaults:
   - training_hyperparams: cityscapes_default_train_params
-  #  - dataset_params: cityscapes_stdc_seg50_dataset_params # TODO: uncomment after DatasetInterface refactor
+  - dataset_params: cityscapes_stdc_seg50_dataset_params # TODO: uncomment after DatasetInterface refactor
   - checkpoint_params: default_checkpoint_params
   - _self_
 
-# ------------------------------------- legacy dataset params ------------------------------------- #
-dataset_params:
-  _convert_: all
-  batch_size: 16
-  val_batch_size: 16
-  dataset_dir: /data/cityscapes
-  crop_size: [ 1024, 512 ]
-  img_size: 1024
-  train_loader_drop_last: True
-  color_jitter: 0.5
-  random_scales: [ 0.125, 1.5 ]
-  eval_scale: 0.5
-  cityscapes_ignored_label: 19
-
-  image_mask_transforms_aug:
-    Compose:
-      transforms:
-        - ColorJitterSeg:
-            brightness: ${dataset_params.color_jitter}
-            contrast: ${dataset_params.color_jitter}
-            saturation: ${dataset_params.color_jitter}
-
-        - RandomFlipSeg
-
-        - RandomRescaleSeg:
-            scales: ${dataset_params.random_scales}
-
-        - PadShortToCropSizeSeg:
-            crop_size: ${dataset_params.crop_size}
-            fill_mask: ${dataset_params.cityscapes_ignored_label}
-
-        - CropImageAndMaskSeg:
-            crop_size: ${dataset_params.crop_size}
-            mode: random
-
-  image_mask_transforms:
-    Compose:
-      transforms:
-        - RescaleSeg:
-            scale_factor: ${dataset_params.eval_scale}
-
-dataset_interface:
-  cityscapes:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 10
-# ------------------------------------- legacy dataset params ------------------------------------- #
-
+train_dataloader: cityscapes_train
+val_dataloader: cityscapes_val
 
 architecture: stdc1_seg
 
-
 arch_params:
   num_classes: 19
   use_aux_heads: True

@@ -42,57 +42,12 @@
 
 defaults:
   - training_hyperparams: cityscapes_default_train_params
-  #  - dataset_params: cityscapes_stdc_seg75_dataset_params # TODO: uncomment after DatasetInterface refactor
+  - dataset_params: cityscapes_stdc_seg75_dataset_params
   - checkpoint_params: default_checkpoint_params
   - _self_
 
-# ------------------------------------- legacy dataset params ------------------------------------- #
-dataset_params:
-  _convert_: all
-  batch_size: 4
-  val_batch_size: 4
-  dataset_dir: /data/cityscapes
-  crop_size: [ 1536, 768 ]
-  img_size: 1024
-  train_loader_drop_last: True
-  color_jitter: 0.5
-  random_scales: [ 0.125, 1.5 ]
-  eval_scale: 0.75
-  cityscapes_ignored_label: 19
-
-  image_mask_transforms_aug:
-    Compose:
-      transforms:
-        - ColorJitterSeg:
-            brightness: ${dataset_params.color_jitter}
-            contrast: ${dataset_params.color_jitter}
-            saturation: ${dataset_params.color_jitter}
-
-        - RandomFlipSeg
-
-        - RandomRescaleSeg:
-            scales: ${dataset_params.random_scales}
-
-        - PadShortToCropSizeSeg:
-            crop_size: ${dataset_params.crop_size}
-            fill_mask: ${dataset_params.cityscapes_ignored_label}
-
-        - CropImageAndMaskSeg:
-            crop_size: ${dataset_params.crop_size}
-            mode: random
-
-  image_mask_transforms:
-    Compose:
-      transforms:
-        - RescaleSeg:
-            scale_factor: ${dataset_params.eval_scale}
-
-dataset_interface:
-  cityscapes:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 10
-# ------------------------------------- legacy dataset params ------------------------------------- #
+train_dataloader: cityscapes_train
+val_dataloader: cityscapes_val
 
 
 architecture: stdc1_seg

@@ -28,6 +28,9 @@ defaults:
   - checkpoint_params: default_checkpoint_params
   - anchors: ssd_anchors
 
+train_dataloader: coco2017_train
+val_dataloader: coco2017_val
+
 architecture: ssd_lite_mobilenet_v2
 
 data_loader_num_workers: 8
@@ -43,9 +46,6 @@ arch_params:
   num_classes: 80
   anchors: ${dboxes}
 
-dataset_interface:
-  coco2017_detection:
-    dataset_params: ${dataset_params}
 resume: False
 training_hyperparams:
   resume: ${resume}

@@ -28,12 +28,8 @@ defaults:
   - arch_params: yolox_s_arch_params
   - checkpoint_params: default_checkpoint_params
 
-dataset_interface:
-  coco2017_detection:
-    dataset_params: ${dataset_params}
-
-
-data_loader_num_workers: 8
+train_dataloader: coco2017_train
+val_dataloader: coco2017_val
 
 model_checkpoints_location: local
 
@@ -46,7 +42,7 @@ architecture: yolox_s
 
 multi_gpu: DDP
 
-experiment_suffix: res${dataset_params.train_image_size}
+experiment_suffix: res${dataset_params.train_dataset_params.input_dim}
 experiment_name: ${architecture}_coco2017_${experiment_suffix}
 
 ckpt_root_dir:

@@ -16,31 +16,13 @@
 
 defaults:
   - training_hyperparams: coco_segmentation_shelfnet_lw_train_params
-  #  - dataset_params: coco_segmentation_dataset_params # TODO: uncomment after DatasetInterface refactor
+  - dataset_params: coco_segmentation_dataset_params
   - arch_params: shelfnet34_lw_arch_params
   - checkpoint_params: default_checkpoint_params
   - _self_
 
-# ------------------------------------- legacy dataset params ------------------------------------- #
-dataset_params:
-  batch_size: 8
-  val_batch_size: 24
-  dataset_dir: "/data/coco/"
-  img_size: 608
-  crop_size: 512
-  train_loader_drop_last: True
-
-sub_classes:
-  _target_: super_gradients.training.utils.segmentation_utils.coco_sub_classes_inclusion_tuples_list
-
-dataset_interface:
-  _target_: super_gradients.training.datasets.dataset_interfaces.dataset_interface.CoCoSegmentationDatasetInterface
-  dataset_params: ${dataset_params}
-  dataset_classes_inclusion_tuples_list: ${sub_classes}
-
-data_loader_num_workers: 8
-# ------------------------------------- legacy dataset params ------------------------------------- #
-
+train_dataloader: coco_segmentation_train
+val_dataloader: coco_segmentation_val
 
 checkpoint_params:
   strict_load: True

@@ -1,8 +1,3 @@
-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

@@ -4,7 +4,7 @@ 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
+  root: ./data/cifar10
   train: True
   transforms:
     - RandomCrop:
@@ -31,7 +31,7 @@ train_dataloader_params:
   pin_memory: True
 
 val_dataset_params:
-  root: /data/cifar10
+  root: ./data/cifar10
   train: False
   transforms:
     - ToTensor

@@ -1,81 +1,3 @@
-data_dir: /data/coco # root path to coco data
-train_subdir: images/train2017 # sub directory path of data_dir containing the train data.
-val_subdir: images/val2017 # sub directory path of data_dir containing the validation data.
-train_json_file: instances_train2017.json # path to coco train json file, data_dir/annotations/train_json_file.
-val_json_file: instances_val2017.json # path to coco validation json file, data_dir/annotations/val_json_file.
-
-cache_dir: # path to a directory that will be used for caching (with numpy.memmap).
-cache_train_images: False
-cache_val_images: False
-
-batch_size: 16 # batch size for trainset
-val_batch_size: 64 # batch size for valset
-train_image_size: 640
-val_image_size: 640
-train_input_dim:
-  - ${dataset_params.train_image_size}
-  - ${dataset_params.train_image_size}
-val_input_dim:
-  - ${dataset_params.val_image_size}
-  - ${dataset_params.val_image_size}
-
-filter_box_candidates: False
-targets_format:
-  _target_: super_gradients.training.utils.detection_utils.DetectionTargetsFormat # targets format
-  value: LABEL_CXCYWH
-
-tight_box_rotation: False
-train_transforms:
-  - _target_: super_gradients.training.transforms.transforms.DetectionMosaic
-    input_dim: ${dataset_params.train_input_dim}
-    prob: 1.
-  - _target_: super_gradients.training.transforms.transforms.DetectionRandomAffine
-    degrees: 10.                  # rotation degrees, randomly sampled from [-degrees, degrees]
-    translate: 0.1                # image translation fraction
-    scales: [0.1, 2]              # random rescale range (keeps size by padding/cropping) after mosaic transform.
-    shear: 2.0                    # shear degrees, randomly sampled from [-degrees, degrees]
-    target_size: ${dataset_params.train_input_dim}
-    filter_box_candidates: False  # whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.
-    wh_thr: 2                     # edge size threshold when filter_box_candidates = True (pixels)
-    area_thr: 0.1                 # threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True
-    ar_thr: 20                    # aspect ratio threshold when filter_box_candidates = True
-  - _target_: super_gradients.training.transforms.transforms.DetectionMixup
-    input_dim: ${dataset_params.train_input_dim}
-    mixup_scale: [0.5, 1.5]         # random rescale range for the additional sample in mixup
-    prob: 1.0                       # probability to apply per-sample mixup
-    flip_prob: 0.5                  # probability to apply horizontal flip
-  - _target_: super_gradients.training.transforms.transforms.DetectionHSV
-    prob: 1.0                       # probability to apply HSV transform
-    hgain: 5                        # HSV transform hue gain (randomly sampled from [-hgain, hgain])
-    sgain: 30                       # HSV transform saturation gain (randomly sampled from [-sgain, sgain])
-    vgain: 30                       # HSV transform value gain (randomly sampled from [-vgain, vgain])
-  - _target_: super_gradients.training.transforms.transforms.DetectionHorizontalFlip
-    prob: 0.5                       # probability to apply horizontal flip
-  - _target_: super_gradients.training.transforms.transforms.DetectionPaddedRescale
-    input_dim: ${dataset_params.train_input_dim}
-    max_targets: 120
-  - _target_: super_gradients.training.transforms.transforms.DetectionTargetsFormatTransform
-    output_format: ${dataset_params.targets_format}
-
-val_transforms:
-  - _target_: super_gradients.training.transforms.transforms.DetectionPaddedRescale
-    input_dim: ${dataset_params.val_input_dim}
-  - _target_: super_gradients.training.transforms.transforms.DetectionTargetsFormatTransform
-    max_targets: 50
-    output_format: ${dataset_params.targets_format}
-
-val_collate_fn: # collate function for valset
-  _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
-train_collate_fn: # collate function for trainset
-  _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
-
-class_inclusion_list: # If not None,every class not included will be ignored.
-train_max_num_samples: # If not None, only specified number of samples will be loaded in train dataset
-val_max_num_samples:   # If not None, only specified number of samples will be loaded in test dataset
-with_crowd: False     # Whether to return "crowd" labels in validation
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-
 train_dataset_params:
   data_dir: /data/coco # root path to coco data
   subdir: images/train2017 # sub directory path of data_dir containing the train data.

@@ -5,7 +5,6 @@ cache_dir: # path to a directory that will be used for caching (with numpy.memma
 cache_train_images: False
 cache_val_images: False
 
-
 batch_size: 32
 val_batch_size: 16
 train_image_size: 320
@@ -126,7 +125,7 @@ val_dataset_params:
   cache: False
   transforms:
     - DetectionPaddedRescale:
-      input_dim: ${dataset_params.val_dataset_params.input_dim}
+        input_dim: ${dataset_params.val_dataset_params.input_dim}
     - DetectionTargetsFormatTransform:
         max_targets: 50
         output_format:

@@ -1,25 +1,6 @@
-batch_size: 64 # batch size for trainset in DatasetInterface
-val_batch_size: 200 # batch size for valset in DatasetInterface
-dataset_dir: /data/Imagenet # path to imagenet directory (local)
-traindir: train # dirname inside dataset_dir holding trainset files
-valdir: val # dirname inside dataset_dir holding valset files
-img_mean: [0.485, 0.456, 0.406] # mean for normalization
-img_std: [0.229, 0.224, 0.225] # std for normalization
-crop_size: 224 # crop size (size of net's input)
-resize_size: 256 # loaded image resize size (appplied first among preprocessing transforms)
-color_jitter: 0.0 # color jitter augmentation (applied only to trainset)
-imagenet_pca_aug: 0.0 # imagenet pca augmentation (applied only to trainset)
-train_interpolation: default # interpolation mode
-rand_augment_config_string: # randaugment config string (see super_gradients/training/datasets/auto_augment.py)
-random_erase_prob: 0.0 # random erase probability (applied only to trainset)
-aug_repeat_count: 0 # amount of repetitions (each repetition of an example is augmented differently) for a trainset example.
-
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
 # Base recipe for ImageNet Datasets amd Dataloaders.
-#img_mean: [0.485, 0.456, 0.406] # mean for normalization
-#img_std: [0.229, 0.224, 0.225]  # std  for normalization
+img_mean: [0.485, 0.456, 0.406] # mean for normalization
+img_std: [0.229, 0.224, 0.225]  # std  for normalization
 
 train_dataset_params:
   root: /data/Imagenet/train

@@ -1,18 +1,6 @@
 defaults:
   - imagenet_dataset_params
 
-batch_size: 64
-color_jitter: 0.4
-random_erase_prob: 0.2
-random_erase_value: random
-train_interpolation: random
-auto_augment_config_string: rand-m9-mstd0.5
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-
 train_dataset_params:
   root: /data/Imagenet/train
   transforms:

@@ -1,19 +1,6 @@
 defaults:
   - imagenet_dataset_params
 
-train_loader_drop_last: True
-batch_size: 256
-val_batch_size: 256
-random_erase_prob: 0.2
-random_erase_value: random
-train_interpolation: random
-config_string: rand-m9-mstd0.5
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-
 train_dataset_params:
   root: /data/Imagenet/train
   transforms:

@@ -1,13 +1,6 @@
 defaults:
   - imagenet_dataset_params
 
-batch_size: 128
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-
 train_dataloader_params:
   batch_size: 128
   num_workers: 16

@@ -1,18 +1,6 @@
 defaults:
   - imagenet_dataset_params
 
-batch_size: 256
-color_jitter: 0.4
-random_erase_prob: 0.2
-random_erase_value: random
-train_interpolation: random
-auto_augment_config_string: rand-m9-mstd0.5
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-
 train_dataset_params:
   root: /data/Imagenet/train
   transforms:

@@ -1,23 +1,6 @@
 defaults:
   - imagenet_dataset_params
 
-resize_size: 236
-random_erase_prob: 0
-random_erase_value: random
-train_interpolation: random
-config_string: rand-m7-mstd0.5
-cutmix: True
-cutmix_params:
-  mixup_alpha: 0.2
-  cutmix_alpha: 1.0
-  label_smoothing: 0.1
-
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-
 train_dataset_params:
   root: /data/Imagenet/train
   transforms:

@@ -1,31 +1,12 @@
 defaults:
   - imagenet_dataset_params
 
-batch_size: 192
-val_batch_size: 256
-random_erase_prob: 0
-random_erase_value: random
-train_interpolation: random
-config_string: rand-m7-mstd0.5
-cutmix: True
-cutmix_params:
-  mixup_alpha: 0.2
-  cutmix_alpha: 1.0
-  label_smoothing: 0.1
-aug_repeat_count: 3
-
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-#
-#train_transform_args:
-#  interpolation: random
-#  color_jitter: [0.4, 0.4, 0.4]
-#  random_erase_prob: 0.
-#  random_erase_value: random
-#  auto_augment_config_string: rand-m7-mstd0.5
+train_transform_args:
+  interpolation: random
+  color_jitter: [0.4, 0.4, 0.4]
+  random_erase_prob: 0.
+  random_erase_value: random
+  auto_augment_config_string: rand-m7-mstd0.5
 
 train_dataset_params:
   root: /data/Imagenet/train

@@ -1,26 +1,6 @@
 defaults:
   - imagenet_dataset_params
 
-resize_size: 249
-batch_size: 64
-random_erase_prob: 0
-random_erase_value: random
-train_interpolation: random
-config_string: rand-m7-mstd0.5
-cutmix: True
-cutmix_params:
-  mixup_alpha: 0.2
-  cutmix_alpha: 1.0
-  label_smoothing: 0.1
-img_mean: [0.5, 0.5, 0.5]
-img_std: [0.5, 0.5, 0.5]
-
-
-# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE
-# TODO: UNCOMMENT BELOW WHEN ABOVE IS REMOVED
-#defaults:
-#  - imagenet_dataset_params
-
 train_dataset_params:
   root: /data/Imagenet/train
   transforms:

@@ -17,12 +17,9 @@ defaults:
 arch_params:
   num_classes: 1000
 
-dataset_interface:
-  _target_: super_gradients.training.datasets.dataset_interfaces.dataset_interface.ImageNetDatasetInterface
-  dataset_params: ${dataset_params}
-  data_dir: /data/Imagenet
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
-data_loader_num_workers: 8
 
 resume: False
 training_hyperparams:

@@ -15,14 +15,13 @@ defaults:
   - arch_params: mobilenet_v2_arch_params
   - checkpoint_params: default_checkpoint_params
 
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
+
 arch_params:
   num_classes: 1000
   dropout: 0.2
 
-dataset_interface:
-  _target_: super_gradients.training.datasets.dataset_interfaces.dataset_interface.ImageNetDatasetInterface
-  dataset_params: ${dataset_params}
-  data_dir: /data/Imagenet
 
 data_loader_num_workers: 8
 

@@ -5,12 +5,8 @@ defaults:
   - dataset_params: imagenet_mobilenetv3_dataset_params
   - checkpoint_params: default_checkpoint_params
 
-dataset_interface:
-  _target_: super_gradients.training.datasets.dataset_interfaces.dataset_interface.ImageNetDatasetInterface
-  dataset_params: ${dataset_params}
-  data_dir: /data/Imagenet
-
-data_loader_num_workers: 16
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
 model_checkpoints_location: local
 resume: False

@@ -33,12 +33,8 @@ arch_params:
   dropout_prob: 0.5
   droppath_prob: 0.0
 
-dataset_interface:
-  _target_: super_gradients.training.datasets.dataset_interfaces.dataset_interface.ImageNetDatasetInterface
-  dataset_params: ${dataset_params}
-  data_dir: /data/Imagenet
-
-data_loader_num_workers: 8
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
 model_checkpoints_location: local
 load_checkpoint: False

@@ -21,12 +21,8 @@ arch_params:
   num_classes: 1000
   build_residual_branches: True
 
-dataset_interface:
-  _target_: super_gradients.training.datasets.dataset_interfaces.dataset_interface.ImageNetDatasetInterface
-  dataset_params: ${dataset_params}
-  data_dir: /data/Imagenet
-
-data_loader_num_workers: 8
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
 
 model_checkpoints_location: local

@@ -21,11 +21,8 @@ defaults:
 arch_params:
   droppath_prob: 0.05
 
-dataset_interface:
-  imagenet:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 8
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
 model_checkpoints_location: local
 resume: False

@@ -18,6 +18,8 @@ defaults:
   - arch_params: default_arch_params
   - checkpoint_params: default_checkpoint_params
 
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
 resume: False
 training_hyperparams:
@@ -64,12 +66,6 @@ student_checkpoint_params:
   pretrained_weights: # a string describing the dataset of the pretrained weights (for example "imagenent").
 
 
-dataset_interface:
-  imagenet:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 8
-
 model_checkpoints_location: local
 
 

@@ -18,11 +18,8 @@ defaults:
   - arch_params: vit_base_arch_params
   - checkpoint_params: vit_base_imagenet_checkpoint_params
 
-dataset_interface:
-  imagenet:
-    dataset_params: ${dataset_params}
-
-data_loader_num_workers: 8
+train_dataloader: imagenet_train
+val_dataloader: imagenet_val
 
 model_checkpoints_location: local
 

@@ -16,7 +16,6 @@ defaults:
   - imagenet_vit_base
 
 dataset_params:
-  batch_size: 32
   train_dataloader_params:
     batch_size: 32
 

# This recipe is for testing purpose only

defaults:
  - training_hyperparams: cifar10_resnet_train_params
  - arch_params: resnet18_cifar_arch_params
  - checkpoint_params: default_checkpoint_params

dataset_interface:
  classification_test_dataset:
    dataset_params:
      batch_size: 10

data_loader_num_workers: 1

resume: False
training_hyperparams:
  resume: $(resume}

experiment_name: test

model_checkpoints_location: local

architecture: resnet18

@@ -1,7 +1,6 @@
 # PACKAGE IMPORTS FOR EXTERNAL USAGE
 import super_gradients.training.utils.distributed_training_utils as distributed_training_utils
-from super_gradients.training.datasets import datasets_utils, DataAugmentation, TestDatasetInterface, SegmentationTestDatasetInterface, \
-    DetectionTestDatasetInterface, ClassificationTestDatasetInterface
+from super_gradients.training.datasets import datasets_utils, DataAugmentation
 from super_gradients.training.models import ARCHITECTURES
 from super_gradients.training.sg_trainer import Trainer
 from super_gradients.training.kd_trainer import KDTrainer
@@ -9,6 +8,5 @@ from super_gradients.training.sg_model import SgModel
 from super_gradients.training.kd_model import KDModel
 from super_gradients.common import MultiGPUMode, StrictLoad, EvaluationType
 
-__all__ = ['distributed_training_utils', 'datasets_utils', 'DataAugmentation', 'TestDatasetInterface',
-           'ARCHITECTURES', 'Trainer', 'KDTrainer', 'MultiGPUMode', 'TestDatasetInterface', 'SegmentationTestDatasetInterface', 'DetectionTestDatasetInterface',
-           'ClassificationTestDatasetInterface', 'StrictLoad', 'SgModel', 'EvaluationType', 'KDModel']
+__all__ = ['distributed_training_utils', 'datasets_utils', 'DataAugmentation',
+           'ARCHITECTURES', 'Trainer', 'KDTrainer', 'MultiGPUMode', 'StrictLoad', 'SgModel', 'EvaluationType', 'KDModel']

@@ -0,0 +1,31 @@
+from .dataloaders import coco2017_train, coco2017_val, coco2017_train_yolox, coco2017_val_yolox, \
+    coco2017_train_ssd_lite_mobilenet_v2, coco2017_val_ssd_lite_mobilenet_v2, imagenet_train, imagenet_val, \
+    imagenet_efficientnet_train, imagenet_efficientnet_val, imagenet_mobilenetv2_train, imagenet_mobilenetv2_val, \
+    imagenet_mobilenetv3_train, imagenet_mobilenetv3_val, imagenet_regnetY_train, imagenet_regnetY_val, \
+    imagenet_resnet50_train, imagenet_resnet50_val, imagenet_resnet50_kd_train, imagenet_resnet50_kd_val, \
+    imagenet_vit_base_train, imagenet_vit_base_val, tiny_imagenet_train, tiny_imagenet_val, cifar10_train, cifar10_val, \
+    cifar100_train, cifar100_val, cityscapes_train, cityscapes_val, cityscapes_stdc_seg50_train, \
+    cityscapes_stdc_seg50_val, cityscapes_stdc_seg75_train, cityscapes_stdc_seg75_val, cityscapes_regseg48_train, \
+    cityscapes_regseg48_val, cityscapes_ddrnet_train, cityscapes_ddrnet_val, coco_segmentation_train, \
+    coco_segmentation_val, pascal_aug_segmentation_train, pascal_aug_segmentation_val, pascal_voc_segmentation_train, \
+    pascal_voc_segmentation_val, supervisely_persons_train, supervisely_persons_val, pascal_voc_detection_train, \
+    pascal_voc_detection_val, get_data_loader, get
+
+__all__ = ["coco2017_train", "coco2017_val", "coco2017_train_yolox", "coco2017_val_yolox",
+           "coco2017_train_ssd_lite_mobilenet_v2", "coco2017_val_ssd_lite_mobilenet_v2", "imagenet_train",
+           "imagenet_val",
+           "imagenet_efficientnet_train", "imagenet_efficientnet_val", "imagenet_mobilenetv2_train",
+           "imagenet_mobilenetv2_val",
+           "imagenet_mobilenetv3_train", "imagenet_mobilenetv3_val", "imagenet_regnetY_train", "imagenet_regnetY_val",
+           "imagenet_resnet50_train", "imagenet_resnet50_val", "imagenet_resnet50_kd_train", "imagenet_resnet50_kd_val",
+           "imagenet_vit_base_train", "imagenet_vit_base_val", "tiny_imagenet_train", "tiny_imagenet_val",
+           "cifar10_train", "cifar10_val",
+           "cifar100_train", "cifar100_val", "cityscapes_train", "cityscapes_val", "cityscapes_stdc_seg50_train",
+           "cityscapes_stdc_seg50_val", "cityscapes_stdc_seg75_train", "cityscapes_stdc_seg75_val",
+           "cityscapes_regseg48_train",
+           "cityscapes_regseg48_val", "cityscapes_ddrnet_train", "cityscapes_ddrnet_val", "coco_segmentation_train",
+           "coco_segmentation_val", "pascal_aug_segmentation_train", "pascal_aug_segmentation_val",
+           "pascal_voc_segmentation_train",
+           "pascal_voc_segmentation_val", "supervisely_persons_train", "supervisely_persons_val",
+           "pascal_voc_detection_train",
+           "pascal_voc_detection_val", "get_data_loader", "get"]

@@ -29,7 +29,7 @@ from super_gradients.training.utils.utils import override_default_params_without
 logger = get_logger(__name__)
 
 
-def get_data_loader(config_name, dataset_cls, train, dataset_params={}, dataloader_params={}):
+def get_data_loader(config_name, dataset_cls, train, dataset_params=None, dataloader_params=None):
     """
     Class for creating dataloaders for taking defaults from yaml files in src/super_gradients/recipes.
 
@@ -44,6 +44,11 @@ def get_data_loader(config_name, dataset_cls, train, dataset_params={}, dataload
     :param dataloader_params: DataLoader params that override the yaml configured defaults, then passed to the DataLoader.__init__
     :return: DataLoader
     """
+    if dataloader_params is None:
+        dataloader_params = dict()
+    if dataset_params is None:
+        dataset_params = dict()
+
     GlobalHydra.instance().clear()
     with initialize_config_dir(config_dir=pkg_resources.resource_filename("super_gradients.recipes", "")):
         # config is relative to a module
@@ -100,7 +105,7 @@ def _instantiate_sampler(dataset, dataloader_params):
     return dataloader_params
 
 
-def coco2017_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco2017_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="coco_detection_dataset_params",
                            dataset_cls=COCODetectionDataset,
                            train=True,
@@ -109,7 +114,7 @@ def coco2017_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def coco2017_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco2017_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="coco_detection_dataset_params",
                            dataset_cls=COCODetectionDataset,
                            train=False,
@@ -118,15 +123,15 @@ def coco2017_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def coco2017_train_yolox(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco2017_train_yolox(dataset_params: Dict = None, dataloader_params: Dict = None):
     return coco2017_train(dataset_params, dataloader_params)
 
 
-def coco2017_val_yolox(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco2017_val_yolox(dataset_params: Dict = None, dataloader_params: Dict = None):
     return coco2017_val(dataset_params, dataloader_params)
 
 
-def coco2017_train_ssd_lite_mobilenet_v2(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco2017_train_ssd_lite_mobilenet_v2(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="coco_detection_ssd_lite_mobilenet_v2_dataset_params",
                            dataset_cls=COCODetectionDataset,
                            train=True,
@@ -135,7 +140,7 @@ def coco2017_train_ssd_lite_mobilenet_v2(dataset_params: Dict = {}, dataloader_p
                            )
 
 
-def coco2017_val_ssd_lite_mobilenet_v2(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco2017_val_ssd_lite_mobilenet_v2(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="coco_detection_ssd_lite_mobilenet_v2_dataset_params",
                            dataset_cls=COCODetectionDataset,
                            train=False,
@@ -144,7 +149,7 @@ def coco2017_val_ssd_lite_mobilenet_v2(dataset_params: Dict = {}, dataloader_par
                            )
 
 
-def imagenet_train(dataset_params={}, dataloader_params={}, config_name="imagenet_dataset_params"):
+def imagenet_train(dataset_params=None, dataloader_params=None, config_name="imagenet_dataset_params"):
     return get_data_loader(config_name=config_name,
                            dataset_cls=ImageNetDataset,
                            train=True,
@@ -152,7 +157,7 @@ def imagenet_train(dataset_params={}, dataloader_params={}, config_name="imagene
                            dataloader_params=dataloader_params)
 
 
-def imagenet_val(dataset_params={}, dataloader_params={}, config_name="imagenet_dataset_params"):
+def imagenet_val(dataset_params=None, dataloader_params=None, config_name="imagenet_dataset_params"):
     return get_data_loader(config_name=config_name,
                            dataset_cls=ImageNetDataset,
                            train=False,
@@ -160,63 +165,63 @@ def imagenet_val(dataset_params={}, dataloader_params={}, config_name="imagenet_
                            dataloader_params=dataloader_params)
 
 
-def imagenet_efficientnet_train(dataset_params={}, dataloader_params={}):
+def imagenet_efficientnet_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_efficientnet_dataset_params")
 
 
-def imagenet_efficientnet_val(dataset_params={}, dataloader_params={}):
+def imagenet_efficientnet_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_efficientnet_dataset_params")
 
 
-def imagenet_mobilenetv2_train(dataset_params={}, dataloader_params={}):
+def imagenet_mobilenetv2_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_mobilenetv2_dataset_params")
 
 
-def imagenet_mobilenetv2_val(dataset_params={}, dataloader_params={}):
+def imagenet_mobilenetv2_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_mobilenetv2_dataset_params")
 
 
-def imagenet_mobilenetv3_train(dataset_params={}, dataloader_params={}):
+def imagenet_mobilenetv3_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_mobilenetv3_dataset_params")
 
 
-def imagenet_mobilenetv3_val(dataset_params={}, dataloader_params={}):
+def imagenet_mobilenetv3_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_mobilenetv3_dataset_params")
 
 
-def imagenet_regnetY_train(dataset_params={}, dataloader_params={}):
+def imagenet_regnetY_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_regnetY_dataset_params")
 
 
-def imagenet_regnetY_val(dataset_params={}, dataloader_params={}):
+def imagenet_regnetY_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_regnetY_dataset_params")
 
 
-def imagenet_resnet50_train(dataset_params={}, dataloader_params={}):
+def imagenet_resnet50_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_resnet50_dataset_params")
 
 
-def imagenet_resnet50_val(dataset_params={}, dataloader_params={}):
+def imagenet_resnet50_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_resnet50_dataset_params")
 
 
-def imagenet_resnet50_kd_train(dataset_params={}, dataloader_params={}):
+def imagenet_resnet50_kd_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_resnet50_kd_dataset_params")
 
 
-def imagenet_resnet50_kd_val(dataset_params={}, dataloader_params={}):
+def imagenet_resnet50_kd_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_resnet50_kd_dataset_params")
 
 
-def imagenet_vit_base_train(dataset_params={}, dataloader_params={}):
+def imagenet_vit_base_train(dataset_params=None, dataloader_params=None):
     return imagenet_train(dataset_params, dataloader_params, config_name="imagenet_vit_base_dataset_params")
 
 
-def imagenet_vit_base_val(dataset_params={}, dataloader_params={}):
+def imagenet_vit_base_val(dataset_params=None, dataloader_params=None):
     return imagenet_val(dataset_params, dataloader_params, config_name="imagenet_vit_base_dataset_params")
 
 
-def tiny_imagenet_train(dataset_params={}, dataloader_params={}, config_name="tiny_imagenet_dataset_params"):
+def tiny_imagenet_train(dataset_params=None, dataloader_params=None, config_name="tiny_imagenet_dataset_params"):
     return get_data_loader(config_name=config_name,
                            dataset_cls=ImageNetDataset,
                            train=True,
@@ -224,7 +229,7 @@ def tiny_imagenet_train(dataset_params={}, dataloader_params={}, config_name="ti
                            dataloader_params=dataloader_params)
 
 
-def tiny_imagenet_val(dataset_params={}, dataloader_params={}, config_name="tiny_imagenet_dataset_params"):
+def tiny_imagenet_val(dataset_params=None, dataloader_params=None, config_name="tiny_imagenet_dataset_params"):
     return get_data_loader(config_name=config_name,
                            dataset_cls=ImageNetDataset,
                            train=False,
@@ -232,7 +237,7 @@ def tiny_imagenet_val(dataset_params={}, dataloader_params={}, config_name="tiny
                            dataloader_params=dataloader_params)
 
 
-def cifar10_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cifar10_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cifar10_dataset_params",
                            dataset_cls=Cifar10,
                            train=True,
@@ -241,7 +246,7 @@ def cifar10_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def cifar10_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cifar10_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cifar10_dataset_params",
                            dataset_cls=Cifar10,
                            train=False,
@@ -250,7 +255,7 @@ def cifar10_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def cifar100_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cifar100_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cifar100_dataset_params",
                            dataset_cls=Cifar100,
                            train=True,
@@ -259,7 +264,7 @@ def cifar100_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def cifar100_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cifar100_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cifar100_dataset_params",
                            dataset_cls=Cifar100,
                            train=False,
@@ -268,28 +273,31 @@ def cifar100_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-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)))
+def classification_test_dataloader(batch_size: int = 5, image_size: int = 32, dataset_size=None) -> DataLoader:
+    dataset_size = dataset_size or batch_size
+    images = torch.Tensor(np.zeros((dataset_size, 3, image_size, image_size)))
+    ground_truth = torch.LongTensor(np.zeros((dataset_size)))
     dataset = TensorDataset(images, ground_truth)
     return DataLoader(dataset=dataset, batch_size=batch_size)
 
 
-def detection_test_dataloader(batch_size: int = 5, image_size: int = 320) -> DataLoader:
-    images = torch.Tensor(np.zeros((batch_size, 3, image_size, image_size)))
-    ground_truth = torch.LongTensor(np.zeros((batch_size, 6)))
+def detection_test_dataloader(batch_size: int = 5, image_size: int = 320, dataset_size=None) -> DataLoader:
+    dataset_size = dataset_size or batch_size
+    images = torch.Tensor(np.zeros((dataset_size, 3, image_size, image_size)))
+    ground_truth = torch.Tensor(np.zeros((dataset_size, 6)))
     dataset = TensorDataset(images, ground_truth)
     return DataLoader(dataset=dataset, batch_size=batch_size)
 
 
-def segmentation_test_dataloader(batch_size: int = 5, image_size: int = 512) -> DataLoader:
-    images = torch.Tensor(np.zeros((batch_size, 3, image_size, image_size)))
-    ground_truth = torch.LongTensor(np.zeros((batch_size, image_size, image_size)))
+def segmentation_test_dataloader(batch_size: int = 5, image_size: int = 512, dataset_size=None) -> DataLoader:
+    dataset_size = dataset_size or batch_size
+    images = torch.Tensor(np.zeros((dataset_size, 3, image_size, image_size)))
+    ground_truth = torch.LongTensor(np.zeros((dataset_size, image_size, image_size)))
     dataset = TensorDataset(images, ground_truth)
     return DataLoader(dataset=dataset, batch_size=batch_size)
 
 
-def cityscapes_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=True,
@@ -298,7 +306,7 @@ def cityscapes_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def cityscapes_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=False,
@@ -307,7 +315,7 @@ def cityscapes_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
                            )
 
 
-def cityscapes_stdc_seg50_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_stdc_seg50_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_stdc_seg50_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=True,
@@ -316,7 +324,7 @@ def cityscapes_stdc_seg50_train(dataset_params: Dict = {}, dataloader_params: Di
                            )
 
 
-def cityscapes_stdc_seg50_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_stdc_seg50_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_stdc_seg50_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=False,
@@ -325,7 +333,7 @@ def cityscapes_stdc_seg50_val(dataset_params: Dict = {}, dataloader_params: Dict
                            )
 
 
-def cityscapes_stdc_seg75_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_stdc_seg75_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_stdc_seg75_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=True,
@@ -334,7 +342,7 @@ def cityscapes_stdc_seg75_train(dataset_params: Dict = {}, dataloader_params: Di
                            )
 
 
-def cityscapes_stdc_seg75_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_stdc_seg75_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_stdc_seg75_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=False,
@@ -343,7 +351,7 @@ def cityscapes_stdc_seg75_val(dataset_params: Dict = {}, dataloader_params: Dict
                            )
 
 
-def cityscapes_regseg48_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_regseg48_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_regseg48_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=True,
@@ -352,7 +360,7 @@ def cityscapes_regseg48_train(dataset_params: Dict = {}, dataloader_params: Dict
                            )
 
 
-def cityscapes_regseg48_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_regseg48_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_regseg48_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=False,
@@ -361,7 +369,7 @@ def cityscapes_regseg48_val(dataset_params: Dict = {}, dataloader_params: Dict =
                            )
 
 
-def cityscapes_ddrnet_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_ddrnet_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_ddrnet_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=True,
@@ -370,7 +378,7 @@ def cityscapes_ddrnet_train(dataset_params: Dict = {}, dataloader_params: Dict =
                            )
 
 
-def cityscapes_ddrnet_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def cityscapes_ddrnet_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="cityscapes_ddrnet_dataset_params",
                            dataset_cls=CityscapesDataset,
                            train=False,
@@ -379,7 +387,7 @@ def cityscapes_ddrnet_val(dataset_params: Dict = {}, dataloader_params: Dict = {
                            )
 
 
-def coco_segmentation_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco_segmentation_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="coco_segmentation_dataset_params",
                            dataset_cls=CoCoSegmentationDataSet,
                            train=True,
@@ -388,7 +396,7 @@ def coco_segmentation_train(dataset_params: Dict = {}, dataloader_params: Dict =
                            )
 
 
-def coco_segmentation_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def coco_segmentation_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="coco_segmentation_dataset_params",
                            dataset_cls=CoCoSegmentationDataSet,
                            train=False,
@@ -397,7 +405,7 @@ def coco_segmentation_val(dataset_params: Dict = {}, dataloader_params: Dict = {
                            )
 
 
-def pascal_aug_segmentation_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def pascal_aug_segmentation_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="pascal_aug_segmentation_dataset_params",
                            dataset_cls=PascalAUG2012SegmentationDataSet,
                            train=True,
@@ -406,7 +414,7 @@ def pascal_aug_segmentation_train(dataset_params: Dict = {}, dataloader_params:
                            )
 
 
-def pascal_aug_segmentation_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def pascal_aug_segmentation_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="pascal_aug_segmentation_dataset_params",
                            dataset_cls=PascalAUG2012SegmentationDataSet,
                            train=False,
@@ -415,7 +423,7 @@ def pascal_aug_segmentation_val(dataset_params: Dict = {}, dataloader_params: Di
                            )
 
 
-def pascal_voc_segmentation_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def pascal_voc_segmentation_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="pascal_voc_segmentation_dataset_params",
                            dataset_cls=PascalVOC2012SegmentationDataSet,
                            train=True,
@@ -424,7 +432,7 @@ def pascal_voc_segmentation_train(dataset_params: Dict = {}, dataloader_params:
                            )
 
 
-def pascal_voc_segmentation_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def pascal_voc_segmentation_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="pascal_voc_segmentation_dataset_params",
                            dataset_cls=PascalVOC2012SegmentationDataSet,
                            train=False,
@@ -433,7 +441,7 @@ def pascal_voc_segmentation_val(dataset_params: Dict = {}, dataloader_params: Di
                            )
 
 
-def supervisely_persons_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def supervisely_persons_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="supervisely_persons_dataset_params",
                            dataset_cls=SuperviselyPersonsDataset,
                            train=True,
@@ -441,7 +449,7 @@ def supervisely_persons_train(dataset_params: Dict = {}, dataloader_params: Dict
                            dataloader_params=dataloader_params)
 
 
-def supervisely_persons_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def supervisely_persons_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="supervisely_persons_dataset_params",
                            dataset_cls=SuperviselyPersonsDataset,
                            train=False,
@@ -449,7 +457,7 @@ def supervisely_persons_val(dataset_params: Dict = {}, dataloader_params: Dict =
                            dataloader_params=dataloader_params)
 
 
-def pascal_voc_detection_train(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def pascal_voc_detection_train(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="pascal_voc_detection_dataset_params",
                            dataset_cls=PascalVOCUnifiedDetectionTrainDataset,
                            train=True,
@@ -458,10 +466,78 @@ def pascal_voc_detection_train(dataset_params: Dict = {}, dataloader_params: Dic
                            )
 
 
-def pascal_voc_detection_val(dataset_params: Dict = {}, dataloader_params: Dict = {}):
+def pascal_voc_detection_val(dataset_params: Dict = None, dataloader_params: Dict = None):
     return get_data_loader(config_name="pascal_voc_detection_dataset_params",
                            dataset_cls=PascalVOCDetectionDataset,
                            train=False,
                            dataset_params=dataset_params,
                            dataloader_params=dataloader_params
                            )
+
+
+ALL_DATALOADERS = {"coco2017_train": coco2017_train,
+                   "coco2017_val": coco2017_val,
+                   "coco2017_train_yolox": coco2017_train_yolox,
+                   "coco2017_val_yolox": coco2017_val_yolox,
+                   "coco2017_train_ssd_lite_mobilenet_v2": coco2017_train_ssd_lite_mobilenet_v2,
+                   "coco2017_val_ssd_lite_mobilenet_v2": coco2017_val_ssd_lite_mobilenet_v2,
+                   "imagenet_train": imagenet_train,
+                   "imagenet_val": imagenet_val,
+                   "imagenet_efficientnet_train": imagenet_efficientnet_train,
+                   "imagenet_efficientnet_val": imagenet_efficientnet_val,
+                   "imagenet_mobilenetv2_train": imagenet_mobilenetv2_train,
+                   "imagenet_mobilenetv2_val": imagenet_mobilenetv2_val,
+                   "imagenet_mobilenetv3_train": imagenet_mobilenetv3_train,
+                   "imagenet_mobilenetv3_val": imagenet_mobilenetv3_val,
+                   "imagenet_regnetY_train": imagenet_regnetY_train,
+                   "imagenet_regnetY_val": imagenet_regnetY_val,
+                   "imagenet_resnet50_train": imagenet_resnet50_train,
+                   "imagenet_resnet50_val": imagenet_resnet50_val,
+                   "imagenet_resnet50_kd_train": imagenet_resnet50_kd_train,
+                   "imagenet_resnet50_kd_val": imagenet_resnet50_kd_val,
+                   "imagenet_vit_base_train": imagenet_vit_base_train,
+                   "imagenet_vit_base_val": imagenet_vit_base_val,
+                   "tiny_imagenet_train": tiny_imagenet_train,
+                   "tiny_imagenet_val": tiny_imagenet_val,
+                   "cifar10_train": cifar10_train,
+                   "cifar10_val": cifar10_val,
+                   "cifar100_train": cifar100_train,
+                   "cifar100_val": cifar100_val,
+                   "cityscapes_train": cityscapes_train,
+                   "cityscapes_val": cityscapes_val,
+                   "cityscapes_stdc_seg50_train": cityscapes_stdc_seg50_train,
+                   "cityscapes_stdc_seg50_val": cityscapes_stdc_seg50_val,
+                   "cityscapes_stdc_seg75_train": cityscapes_stdc_seg75_train,
+                   "cityscapes_stdc_seg75_val": cityscapes_stdc_seg75_val,
+                   "cityscapes_regseg48_train": cityscapes_regseg48_train,
+                   "cityscapes_regseg48_val": cityscapes_regseg48_val,
+                   "cityscapes_ddrnet_train": cityscapes_ddrnet_train,
+                   "cityscapes_ddrnet_val": cityscapes_ddrnet_val,
+                   "coco_segmentation_train": coco_segmentation_train,
+                   "coco_segmentation_val": coco_segmentation_val,
+                   "pascal_aug_segmentation_train": pascal_aug_segmentation_train,
+                   "pascal_aug_segmentation_val": pascal_aug_segmentation_val,
+                   "pascal_voc_segmentation_train": pascal_voc_segmentation_train,
+                   "pascal_voc_segmentation_val": pascal_voc_segmentation_val,
+                   "supervisely_persons_train": supervisely_persons_train,
+                   "supervisely_persons_val": supervisely_persons_val,
+                   "pascal_voc_detection_train": pascal_voc_detection_train,
+                   "pascal_voc_detection_val": pascal_voc_detection_val
+                   }
+
+
+def get(name: str, dataset_params: Dict = None, dataloader_params: Dict = None):
+    """
+
+    
+    :param name: 
+    :param dataset_params: 
+    :param dataloader_params: 
+    :return: 
+    """
+
+    if name not in ALL_DATALOADERS.keys():
+        raise ValueError("Unsupported dataloader: " + str(name))
+
+    dataloader_cls = ALL_DATALOADERS[name]
+    return dataloader_cls(dataset_params=dataset_params, dataloader_params=dataloader_params)

@@ -3,8 +3,6 @@ import cv2
 
 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, Cifar10, Cifar100
 from super_gradients.training.datasets.detection_datasets import DetectionDataset, COCODetectionDataset, PascalVOCDetectionDataset
 from super_gradients.training.datasets.segmentation_datasets.segmentation_dataset import SegmentationDataSet
@@ -13,23 +11,12 @@ from super_gradients.training.datasets.segmentation_datasets.pascal_aug_segmenta
 from super_gradients.training.datasets.segmentation_datasets.coco_segmentation import CoCoSegmentationDataSet
 from super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation import SuperviselyPersonsDataset
 
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import TestDatasetInterface, DatasetInterface, \
-    Cifar10DatasetInterface, CoCoSegmentationDatasetInterface, \
-    PascalVOC2012SegmentationDataSetInterface, PascalAUG2012SegmentationDataSetInterface, \
-    TestYoloDetectionDatasetInterface, SegmentationTestDatasetInterface, DetectionTestDatasetInterface, \
-    ClassificationTestDatasetInterface, ImageNetDatasetInterface
 
 cv2.setNumThreads(0)
 
 
-__all__ = ['DataAugmentation', 'ListDataset', 'DirectoryDataSet', 'CLASSIFICATION_DATASETS', 'OBJECT_DETECTION_DATASETS',
-           'SEMANTIC_SEGMENTATION_DATASETS', 'SegmentationDataSet',
+__all__ = ['DataAugmentation', 'ListDataset', 'DirectoryDataSet', 'SegmentationDataSet',
            'PascalVOC2012SegmentationDataSet',
-           'PascalAUG2012SegmentationDataSet', 'CoCoSegmentationDataSet', 'TestDatasetInterface', 'DatasetInterface',
-           'Cifar10DatasetInterface', 'CoCoSegmentationDatasetInterface',
-           'PascalVOC2012SegmentationDataSetInterface', 'PascalAUG2012SegmentationDataSetInterface',
-           'TestYoloDetectionDatasetInterface', 'DetectionTestDatasetInterface', 'ClassificationTestDatasetInterface',
-           'SegmentationTestDatasetInterface',
-           'ImageNetDatasetInterface',
+           'PascalAUG2012SegmentationDataSet', 'CoCoSegmentationDataSet',
            'DetectionDataset', 'COCODetectionDataset', 'PascalVOCDetectionDataset', 'ImageNetDataset',
            'Cifar10', 'Cifar100', 'SuperviselyPersonsDataset']

from collections import defaultdict
from typing import Dict, List, Type

from super_gradients.training.datasets.dataset_interfaces import DatasetInterface, TestDatasetInterface, \
    LibraryDatasetInterface, \
    ClassificationDatasetInterface, Cifar10DatasetInterface, Cifar100DatasetInterface, \
    ImageNetDatasetInterface, TinyImageNetDatasetInterface, CoCoSegmentationDatasetInterface,\
    PascalAUG2012SegmentationDataSetInterface, PascalVOC2012SegmentationDataSetInterface
from super_gradients.common.data_types.enum.deep_learning_task import DeepLearningTask
from super_gradients.training.datasets.dataset_interfaces.dataset_interface import CoCoDetectionDatasetInterface

CLASSIFICATION_DATASETS = {
    "test_dataset": TestDatasetInterface,
    "library_dataset": LibraryDatasetInterface,
    "classification_dataset": ClassificationDatasetInterface,
    "cifar_10": Cifar10DatasetInterface,
    "cifar_100": Cifar100DatasetInterface,
    "imagenet": ImageNetDatasetInterface,
    "tiny_imagenet": TinyImageNetDatasetInterface
}

OBJECT_DETECTION_DATASETS = {
    "coco": CoCoDetectionDatasetInterface,
}

SEMANTIC_SEGMENTATION_DATASETS = {
    "coco": CoCoSegmentationDatasetInterface,
    "pascal_voc": PascalVOC2012SegmentationDataSetInterface,
    "pascal_aug": PascalAUG2012SegmentationDataSetInterface
}


class DataSetDoesNotExistException(Exception):
    """
    The requested dataset does not exist, or is not implemented.
    """
    pass


class SgLibraryDatasets(object):
    """
    Holds all of the different library dataset dictionaries, by DL Task mapping

        Attributes:
            CLASSIFICATION          Dictionary of Classification Data sets
            OBJECT_DETECTION        Dictionary of Object Detection Data sets
            SEMANTIC_SEGMENTATION   Dictionary of Semantic Segmentation Data sets
    """
    CLASSIFICATION = CLASSIFICATION_DATASETS
    OBJECT_DETECTION = OBJECT_DETECTION_DATASETS
    SEMANTIC_SEGMENTATION = SEMANTIC_SEGMENTATION_DATASETS

    _datasets_mapping = {
        DeepLearningTask.CLASSIFICATION: CLASSIFICATION,
        DeepLearningTask.SEMANTIC_SEGMENTATION: SEMANTIC_SEGMENTATION,
        DeepLearningTask.OBJECT_DETECTION: OBJECT_DETECTION,
    }

    @staticmethod
    def get_all_available_datasets() -> Dict[str, List[str]]:
        """
        Gets all the available datasets.
        """
        all_datasets: Dict[str, List[str]] = defaultdict(list)
        for dl_task, task_datasets in SgLibraryDatasets._datasets_mapping.items():
            for dataset_name, dataset_interface in task_datasets.items():
                all_datasets[dl_task].append(dataset_name)

        # TODO: Return Dataset Metadata list from the dataset interfaces objects
        # TODO: Transform DatasetInterface -> DataSetMetadata
        return all_datasets

    @staticmethod
    def get_dataset(dl_task: str, dataset_name: str) -> Type[DatasetInterface]:
        """
        Get's a dataset with a given name for a given deep learning task.
        examp:
        >>> SgLibraryDatasets.get_dataset(dl_task='classification', dataset_name='cifar_100')
        >>> <Cifar100DatasetInterface instance>
        """
        task_datasets: Dict[str, DatasetInterface] = SgLibraryDatasets._datasets_mapping.get(dl_task)
        if not task_datasets:
            raise ValueError(f"Invalid Deep Learining Task: {dl_task}")

        dataset: DatasetInterface = task_datasets.get(dataset_name)
        if not dataset:
            raise DataSetDoesNotExistException(dataset_name)

        return dataset

from super_gradients.training.datasets.dataset_interfaces.dataset_interface import DatasetInterface, TestDatasetInterface, \
    LibraryDatasetInterface, \
    ClassificationDatasetInterface, Cifar10DatasetInterface, Cifar100DatasetInterface, \
    ImageNetDatasetInterface, TinyImageNetDatasetInterface, CoCoSegmentationDatasetInterface, \
    PascalAUG2012SegmentationDataSetInterface, PascalVOC2012SegmentationDataSetInterface, \
    TestYoloDetectionDatasetInterface, SegmentationTestDatasetInterface, DetectionTestDatasetInterface, ClassificationTestDatasetInterface,\
    CoCoDetectionDatasetInterface, PascalVOCUnifiedDetectionDatasetInterface


__all__ = ['DatasetInterface', 'TestDatasetInterface', 'LibraryDatasetInterface', 'ClassificationDatasetInterface', 'Cifar10DatasetInterface',
           'Cifar100DatasetInterface', 'ImageNetDatasetInterface', 'TinyImageNetDatasetInterface',
           'CoCoSegmentationDatasetInterface', 'PascalAUG2012SegmentationDataSetInterface',
           'PascalVOC2012SegmentationDataSetInterface', 'TestYoloDetectionDatasetInterface', 'SegmentationTestDatasetInterface',
           'DetectionTestDatasetInterface', 'ClassificationTestDatasetInterface', 'CoCoDetectionDatasetInterface',
           'PascalVOCUnifiedDetectionDatasetInterface']

import os

import numpy as np
import torch
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import ConcatDataset, BatchSampler, DataLoader
from torch.utils.data.distributed import DistributedSampler

from super_gradients.common import DatasetDataInterface
from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.common.environment import AWS_ENV_NAME
from super_gradients.training import utils as core_utils
from super_gradients.training.datasets import datasets_utils, DataAugmentation
from super_gradients.training.datasets.data_augmentation import Lighting, RandomErase
from super_gradients.training.datasets.datasets_conf import COCO_DETECTION_CLASSES_LIST
from super_gradients.training.datasets.datasets_utils import RandomResizedCropAndInterpolation, worker_init_reset_seed
from super_gradients.training.datasets.detection_datasets import COCODetectionDataset, PascalVOCDetectionDataset
from super_gradients.training.datasets.mixup import CollateMixup
from super_gradients.training.datasets.samplers.infinite_sampler import InfiniteSampler
from super_gradients.training.datasets.samplers.repeated_augmentation_sampler import RepeatAugSampler
from super_gradients.training.datasets.segmentation_datasets import PascalVOC2012SegmentationDataSet, \
    PascalAUG2012SegmentationDataSet, CoCoSegmentationDataSet
from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset
from super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation import \
    SuperviselyPersonsDataset
from super_gradients.training.exceptions.dataset_exceptions import IllegalDatasetParameterException
from super_gradients.training.transforms.transforms import RandomFlip, Rescale, RandomRescale, CropImageAndMask, \
    PadShortToCropSize
from super_gradients.training.utils import get_param
from super_gradients.training.utils.distributed_training_utils import get_local_rank, wait_for_the_master

default_dataset_params = {"batch_size": 64, "val_batch_size": 200, "test_batch_size": 200, "dataset_dir": "./data/",
                          "s3_link": None}
LIBRARY_DATASETS = {
    "cifar10": {'class': datasets.CIFAR10, 'mean': (0.4914, 0.4822, 0.4465), 'std': (0.2023, 0.1994, 0.2010)},
    "cifar100": {'class': datasets.CIFAR100, 'mean': (0.5071, 0.4865, 0.4409), 'std': (0.2673, 0.2564, 0.2762)},
    "SVHN": {'class': datasets.SVHN, 'mean': None, 'std': None}
}

logger = get_logger(__name__)


class DatasetInterface:
    """
    DatasetInterface - This class manages all of the "communiation" the Model has with the Data Sets
    """

    def __init__(self, dataset_params={}, train_loader=None, val_loader=None, test_loader=None, classes=None):
        """
        @param train_loader: torch.utils.data.Dataloader (optional) dataloader for training.
        @param test_loader: torch.utils.data.Dataloader (optional) dataloader for testing.
        @param classes: list of classes.

        Note: the above parameters will be discarded in case dataset_params is passed.

        @param dataset_params:

            - `batch_size` : int (default=64)

                Number of examples per batch for training. Large batch sizes are recommended.

            - `val_batch_size` : int (default=200)

                Number of examples per batch for validation. Large batch sizes are recommended.

            - `dataset_dir` : str (default="./data/")

                Directory location for the data. Data will be downloaded to this directory when getting it from a
                remote url.

            - `s3_link` : str (default=None)

                remote s3 link to download the data (optional).

            - `aug_repeat_count` : int (default=0)

                amount of repetitions (each repetition of an example is augmented differently) for each
                 example for the trainset.

        """

        self.dataset_params = core_utils.HpmStruct(**default_dataset_params)
        self.dataset_params.override(**dataset_params)

        self.trainset, self.valset, self.testset = None, None, None
        self.train_loader, self.val_loader, self.test_loader = train_loader, val_loader, test_loader
        self.classes = classes
        self.batch_size_factor = 1
        if self.dataset_params.s3_link is not None:
            self.download_from_cloud()

    def download_from_cloud(self):
        if self.dataset_params.s3_link is not None:
            env_name = AWS_ENV_NAME
            downloader = DatasetDataInterface(env=env_name)
            target_dir = self.dataset_params.dataset_dir
            if not os.path.exists(target_dir):
                os.mkdir(target_dir)
            downloader.load_remote_dataset_file(self.dataset_params.s3_link, target_dir)

    def build_data_loaders(self, batch_size_factor=1, num_workers=8, train_batch_size=None, val_batch_size=None,
                           test_batch_size=None, distributed_sampler: bool = False):
        """

        define train, val (and optionally test) loaders. The method deals separately with distributed training and standard
        (non distributed, or parallel training). In the case of distributed training we need to rely on distributed
        samplers.
        :param batch_size_factor: int - factor to multiply the batch size (usually for multi gpu)
        :param num_workers: int - number of workers (parallel processes) for dataloaders
        :param train_batch_size: int - batch size for train loader, if None will be taken from dataset_params
        :param val_batch_size: int - batch size for val loader, if None will be taken from dataset_params
        :param distributed_sampler: boolean flag for distributed training mode
        :return: train_loader, val_loader, classes: list of classes
        """
        # CHANGE THE BATCH SIZE ACCORDING TO THE NUMBER OF DEVICES - ONLY IN NON-DISTRIBUTED TRAINING MODE
        # IN DISTRIBUTED MODE WE NEED DISTRIBUTED SAMPLERS
        # NO SHUFFLE IN DISTRIBUTED TRAINING

        aug_repeat_count = get_param(self.dataset_params, "aug_repeat_count", 0)
        if aug_repeat_count > 0 and not distributed_sampler:
            raise IllegalDatasetParameterException("repeated augmentation is only supported with DDP.")

        if distributed_sampler:
            self.batch_size_factor = 1
            train_sampler = RepeatAugSampler(self.trainset,
                                             num_repeats=aug_repeat_count) if aug_repeat_count > 0 else DistributedSampler(
                self.trainset)
            val_sampler = DistributedSampler(self.valset)
            test_sampler = DistributedSampler(self.testset) if self.testset is not None else None
            train_shuffle = False
        else:
            self.batch_size_factor = batch_size_factor
            train_sampler = None
            val_sampler = None
            test_sampler = None
            train_shuffle = True

        if train_batch_size is None:
            train_batch_size = self.dataset_params.batch_size * self.batch_size_factor
        if val_batch_size is None:
            val_batch_size = self.dataset_params.val_batch_size * self.batch_size_factor
        if test_batch_size is None:
            test_batch_size = self.dataset_params.test_batch_size * self.batch_size_factor

        train_loader_drop_last = core_utils.get_param(self.dataset_params, 'train_loader_drop_last', default_val=False)

        cutmix = core_utils.get_param(self.dataset_params, 'cutmix', False)
        cutmix_params = core_utils.get_param(self.dataset_params, 'cutmix_params')

        # WRAPPING collate_fn
        train_collate_fn = core_utils.get_param(self.trainset, 'collate_fn')
        val_collate_fn = core_utils.get_param(self.valset, 'collate_fn')
        test_collate_fn = core_utils.get_param(self.testset, 'collate_fn')

        if cutmix and train_collate_fn is not None:
            raise IllegalDatasetParameterException("cutmix and collate function cannot be used together")

        if cutmix:
            # FIXME - cutmix should be available only in classification dataset. once we make sure all classification
            # datasets inherit from the same super class, we should move cutmix code to that class
            logger.warning("Cutmix/mixup was enabled. This feature is currently supported only "
                           "for classification datasets.")
            train_collate_fn = CollateMixup(**cutmix_params)

        # FIXME - UNDERSTAND IF THE num_replicas VARIBALE IS NEEDED
        # train_sampler = DistributedSampler(self.trainset,
        #                                    num_replicas=distributed_gpus_num) if distributed_sampler else None
        # val_sampler = DistributedSampler(self.valset,
        #                                   num_replicas=distributed_gpus_num) if distributed_sampler else None

        self.train_loader = torch.utils.data.DataLoader(self.trainset,
                                                        batch_size=train_batch_size,
                                                        shuffle=train_shuffle,
                                                        num_workers=num_workers,
                                                        pin_memory=True,
                                                        sampler=train_sampler,
                                                        collate_fn=train_collate_fn,
                                                        drop_last=train_loader_drop_last)

        self.val_loader = torch.utils.data.DataLoader(self.valset,
                                                      batch_size=val_batch_size,
                                                      shuffle=False,
                                                      num_workers=num_workers,
                                                      pin_memory=True,
                                                      sampler=val_sampler,
                                                      collate_fn=val_collate_fn)

        if self.testset is not None:
            self.test_loader = torch.utils.data.DataLoader(self.testset,
                                                           batch_size=test_batch_size,
                                                           shuffle=False,
                                                           num_workers=num_workers,
                                                           pin_memory=True,
                                                           sampler=test_sampler,
                                                           collate_fn=test_collate_fn)

        self.classes = self.trainset.classes

    def get_data_loaders(self, **kwargs):
        """
        Get self.train_loader, self.val_loader, self.test_loader, self.classes.

        If the data loaders haven't been initialized yet, build them first.

        :param kwargs: kwargs are passed to build_data_loaders.

        """

        if self.train_loader is None and self.val_loader is None:
            self.build_data_loaders(**kwargs)

        return self.train_loader, self.val_loader, self.test_loader, self.classes

    def get_val_sample(self, num_samples=1):
        if num_samples > len(self.valset):
            raise Exception("Tried to load more samples than val-set size")
        if num_samples == 1:
            return self.valset[0]
        else:
            return self.valset[0:num_samples]

    def get_dataset_params(self):
        return self.dataset_params

    def print_dataset_details(self):
        logger.info("{} training samples, {} val samples, {} classes".format(len(self.trainset), len(self.valset),
                                                                             len(self.trainset.classes)))


class ExternalDatasetInterface(DatasetInterface):
    def __init__(self, train_loader, val_loader, num_classes, dataset_params={}):
        """
        ExternalDatasetInterface - A wrapper for external dataset interface that gets dataloaders from keras/TF
        and converts them to Torch-like dataloaders that return torch.Tensors after
        optional collate_fn while maintaining the same interface (connect_dataset_interface etc.)
            :train_loader:       The external train_loader
            :val_loader:        The external val_loader
            :num_classes:        The number of classes
            :dataset_params      The dict that includes the batch_size and/or the collate_fn

            :return:             DataLoaders that generate torch.Tensors batches after collate_fn
        """
        super().__init__(dataset_params)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.classes = num_classes

    def get_data_loaders(self, batch_size_factor: int = 1, num_workers: int = 8, train_batch_size: int = None,
                         val_batch_size: int = None, distributed_sampler: bool = False):

        # CHANGE THE BATCH SIZE ACCORDING TO THE NUMBER OF DEVICES - ONLY IN NON-DISTRIBUED TRAINING MODE
        # IN DISTRIBUTED MODE WE NEED DISTRIBUTED SAMPLERS
        # NO SHUFFLE IN DISTRIBUTED TRAINING
        if distributed_sampler:
            self.batch_size_factor = 1
            train_sampler = DistributedSampler(self.trainset, shuffle=True)
            val_sampler = DistributedSampler(self.valset)
            train_shuffle = False
        else:
            self.batch_size_factor = batch_size_factor
            train_sampler = None
            val_sampler = None
            train_shuffle = True

        if train_batch_size is None:
            train_batch_size = self.dataset_params.batch_size * self.batch_size_factor
        if val_batch_size is None:
            val_batch_size = self.dataset_params.val_batch_size * self.batch_size_factor

        train_loader_drop_last = core_utils.get_param(self.dataset_params, 'train_loader_drop_last', default_val=False)

        # WRAPPING collate_fn
        train_collate_fn = core_utils.get_param(self.dataset_params, 'train_collate_fn')
        val_collate_fn = core_utils.get_param(self.dataset_params, 'val_collate_fn')

        # FIXME - UNDERSTAND IF THE num_replicas VARIBALE IS NEEDED
        # train_sampler = DistributedSampler(self.trainset,
        #                                    num_replicas=distributed_gpus_num) if distributed_sampler else None
        # val_sampler = DistributedSampler(self.valset,
        #                                   num_replicas=distributed_gpus_num) if distributed_sampler else None

        self.torch_train_loader = torch.utils.data.DataLoader(self.train_loader,
                                                              batch_size=train_batch_size,
                                                              shuffle=train_shuffle,
                                                              num_workers=num_workers,
                                                              pin_memory=True,
                                                              sampler=train_sampler,
                                                              collate_fn=train_collate_fn,
                                                              drop_last=train_loader_drop_last)

        self.torch_val_loader = torch.utils.data.DataLoader(self.val_loader,
                                                            batch_size=val_batch_size,
                                                            shuffle=False,
                                                            num_workers=num_workers,
                                                            pin_memory=True,
                                                            sampler=val_sampler,
                                                            collate_fn=val_collate_fn)

        return self.torch_train_loader, self.torch_val_loader, None, self.classes


class LibraryDatasetInterface(DatasetInterface):
    def __init__(self, name="cifar10", dataset_params={}, to_cutout=False):
        super(LibraryDatasetInterface, self).__init__(dataset_params)
        self.dataset_name = name
        if self.dataset_name not in LIBRARY_DATASETS.keys():
            raise Exception('dataset not found')
        self.lib_dataset_params = LIBRARY_DATASETS[self.dataset_name]

        if self.lib_dataset_params['mean'] is None:
            trainset = torchvision.datasets.SVHN(root=self.dataset_params.dataset_dir, split='train', download=True,
                                                 transform=transforms.ToTensor())
            self.lib_dataset_params['mean'], self.lib_dataset_params['std'] = datasets_utils.get_mean_and_std(trainset)

        # OVERWRITE MEAN AND STD IF DEFINED IN DATASET PARAMS
        self.lib_dataset_params['mean'] = core_utils.get_param(self.dataset_params, 'img_mean',
                                                               default_val=self.lib_dataset_params['mean'])
        self.lib_dataset_params['std'] = core_utils.get_param(self.dataset_params, 'img_std',
                                                              default_val=self.lib_dataset_params['std'])

        crop_size = core_utils.get_param(self.dataset_params, 'crop_size', default_val=32)

        if to_cutout:
            transform_train = transforms.Compose([
                transforms.RandomCrop(crop_size, padding=4),
                transforms.RandomHorizontalFlip(),
                DataAugmentation.normalize(self.lib_dataset_params['mean'], self.lib_dataset_params['std']),
                DataAugmentation.cutout(16),
                DataAugmentation.to_tensor()
            ])
        else:
            transform_train = transforms.Compose([
                transforms.RandomCrop(crop_size, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(self.lib_dataset_params['mean'], self.lib_dataset_params['std']),
            ])

        transform_val = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(self.lib_dataset_params['mean'], self.lib_dataset_params['std']),
        ])
        dataset_cls = self.lib_dataset_params["class"]
        self.trainset = dataset_cls(root=self.dataset_params.dataset_dir, train=True, download=True,
                                    transform=transform_train)

        self.valset = dataset_cls(root=self.dataset_params.dataset_dir, train=False, download=True,
                                  transform=transform_val)


class Cifar10DatasetInterface(LibraryDatasetInterface):
    def __init__(self, dataset_params={}):
        super(Cifar10DatasetInterface, self).__init__(name="cifar10", dataset_params=dataset_params)


class Cifar100DatasetInterface(LibraryDatasetInterface):
    def __init__(self, dataset_params={}):
        super(Cifar100DatasetInterface, self).__init__(name="cifar100", dataset_params=dataset_params)


class TestDatasetInterface(DatasetInterface):
    def __init__(self, trainset, dataset_params={}, classes=None):
        super(TestDatasetInterface, self).__init__(dataset_params)

        self.trainset = trainset
        self.valset = self.trainset
        self.testset = self.trainset
        self.classes = classes

    def get_data_loaders(self, batch_size_factor=1, num_workers=8, train_batch_size=None, val_batch_size=None,
                         distributed_sampler=False):
        self.trainset.classes = [0, 1, 2, 3, 4] if self.classes is None else self.classes
        return super().get_data_loaders(batch_size_factor=batch_size_factor,
                                        num_workers=num_workers,
                                        train_batch_size=train_batch_size,
                                        val_batch_size=val_batch_size,
                                        distributed_sampler=distributed_sampler)


class ClassificationTestDatasetInterface(TestDatasetInterface):
    def __init__(self, dataset_params={}, image_size=32, batch_size=5, classes=None):
        trainset = torch.utils.data.TensorDataset(torch.Tensor(np.zeros((batch_size, 3, image_size, image_size))),
                                                  torch.LongTensor(np.zeros((batch_size))))
        super(ClassificationTestDatasetInterface, self).__init__(trainset=trainset, dataset_params=dataset_params,
                                                                 classes=classes)


class SegmentationTestDatasetInterface(TestDatasetInterface):
    def __init__(self, dataset_params={}, image_size=512, batch_size=4):
        trainset = torch.utils.data.TensorDataset(torch.Tensor(np.zeros((batch_size, 3, image_size, image_size))),
                                                  torch.LongTensor(np.zeros((batch_size, image_size, image_size))))

        super(SegmentationTestDatasetInterface, self).__init__(trainset=trainset, dataset_params=dataset_params)


class DetectionTestDatasetInterface(TestDatasetInterface):
    def __init__(self, dataset_params={}, image_size=320, batch_size=4, classes=None):
        trainset = torch.utils.data.TensorDataset(torch.Tensor(np.zeros((batch_size, 3, image_size, image_size))),
                                                  torch.Tensor(np.zeros((batch_size, 6))))

        super(DetectionTestDatasetInterface, self).__init__(trainset=trainset, dataset_params=dataset_params,
                                                            classes=classes)


class TestYoloDetectionDatasetInterface(DatasetInterface):
    """
    note: the output size is (batch_size, 6) in the test while in real training
    the size of axis 0 can vary (the number of bounding boxes)
    """

    def __init__(self, dataset_params={}, input_dims=(3, 32, 32), batch_size=5):
        super().__init__(dataset_params)
        self.trainset = torch.utils.data.TensorDataset(torch.ones((batch_size, *input_dims)),
                                                       torch.ones((batch_size, 6)))
        self.trainset.classes = [0, 1, 2, 3, 4]
        self.valset = self.trainset


class ImageNetDatasetInterface(DatasetInterface):
    def __init__(self, dataset_params={}, data_dir="/data/Imagenet"):
        super(ImageNetDatasetInterface, self).__init__(dataset_params)

        data_dir = dataset_params['dataset_dir'] if 'dataset_dir' in dataset_params.keys() else data_dir
        traindir = os.path.join(os.path.abspath(data_dir), 'train')
        valdir = os.path.join(data_dir, 'val')
        img_mean = core_utils.get_param(self.dataset_params, 'img_mean', default_val=[0.485, 0.456, 0.406])
        img_std = core_utils.get_param(self.dataset_params, 'img_std', default_val=[0.229, 0.224, 0.225])
        normalize = transforms.Normalize(mean=img_mean, std=img_std)

        crop_size = core_utils.get_param(self.dataset_params, 'crop_size', default_val=224)
        resize_size = core_utils.get_param(self.dataset_params, 'resize_size', default_val=256)
        color_jitter = core_utils.get_param(self.dataset_params, 'color_jitter', default_val=0.0)
        imagenet_pca_aug = core_utils.get_param(self.dataset_params, 'imagenet_pca_aug', default_val=0.0)
        train_interpolation = core_utils.get_param(self.dataset_params, 'train_interpolation', default_val='default')
        rand_augment_config_string = core_utils.get_param(self.dataset_params, 'rand_augment_config_string',
                                                          default_val=None)

        color_jitter = (float(color_jitter),) * 3 if isinstance(color_jitter, float) else color_jitter
        assert len(color_jitter) in (3, 4), "color_jitter must be a scalar or tuple of len 3 or 4"

        color_augmentation = datasets_utils.get_color_augmentation(rand_augment_config_string, color_jitter,
                                                                   crop_size=crop_size, img_mean=img_mean)

        train_transformation_list = [
            RandomResizedCropAndInterpolation(crop_size, interpolation=train_interpolation),
            transforms.RandomHorizontalFlip(),
            color_augmentation,
            transforms.ToTensor(),
            Lighting(imagenet_pca_aug),
            normalize]

        rndm_erase_prob = core_utils.get_param(self.dataset_params, 'random_erase_prob', default_val=0.)
        if rndm_erase_prob:
            train_transformation_list.append(RandomErase(rndm_erase_prob, self.dataset_params.random_erase_value))

        self.trainset = datasets.ImageFolder(traindir, transforms.Compose(train_transformation_list))
        self.valset = datasets.ImageFolder(valdir, transforms.Compose([
            transforms.Resize(resize_size),
            transforms.CenterCrop(crop_size),
            transforms.ToTensor(),
            normalize,
        ]))


class TinyImageNetDatasetInterface(DatasetInterface):
    def __init__(self, dataset_params={}, data_dir="/data/TinyImagenet"):
        super(TinyImageNetDatasetInterface, self).__init__(dataset_params)

        data_dir = dataset_params['dataset_dir'] if 'dataset_dir' in dataset_params.keys() else data_dir
        traindir = os.path.join(os.path.abspath(data_dir), 'train')
        valdir = os.path.join(data_dir, 'val')

        img_mean = core_utils.get_param(self.dataset_params, 'img_mean', default_val=[0.4802, 0.4481, 0.3975])
        img_std = core_utils.get_param(self.dataset_params, 'img_std', default_val=[0.2770, 0.2691, 0.2821])
        normalize = transforms.Normalize(mean=img_mean,
                                         std=img_std)

        crop_size = core_utils.get_param(self.dataset_params, 'crop_size', default_val=56)
        resize_size = core_utils.get_param(self.dataset_params, 'resize_size', default_val=64)

        self.trainset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(crop_size),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))
        self.valset = datasets.ImageFolder(valdir, transforms.Compose([
            transforms.Resize(resize_size),
            transforms.CenterCrop(crop_size),
            transforms.ToTensor(),
            normalize,
        ]))


class ClassificationDatasetInterface(DatasetInterface):
    def __init__(self, normalization_mean=(0, 0, 0), normalization_std=(1, 1, 1), resolution=64,
                 dataset_params={}):
        super(ClassificationDatasetInterface, self).__init__(dataset_params)
        data_dir = self.dataset_params.dataset_dir

        traindir = os.path.join(os.path.abspath(data_dir), 'train')
        valdir = os.path.join(data_dir, 'val')
        normalize = transforms.Normalize(mean=normalization_mean,
                                         std=normalization_std)

        self.trainset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(resolution),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))
        self.valset = datasets.ImageFolder(valdir, transforms.Compose([
            transforms.Resize(int(resolution * 1.15)),
            transforms.CenterCrop(resolution),
            transforms.ToTensor(),
            normalize,
        ]))
        self.data_dir = data_dir
        self.normalization_mean = normalization_mean
        self.normalization_std = normalization_std


class PascalVOC2012SegmentationDataSetInterface(DatasetInterface):
    def __init__(self, dataset_params=None, cache_labels=False, cache_images=False):
        if dataset_params is None:
            dataset_params = dict()
        super().__init__(dataset_params=dataset_params)

        self.root_dir = dataset_params['dataset_dir'] if 'dataset_dir' in dataset_params.keys() \
            else '/data/pascal_voc_2012/VOCdevkit/VOC2012/'

        self.trainset = PascalVOC2012SegmentationDataSet(root=self.root_dir,
                                                         list_file='ImageSets/Segmentation/train.txt',
                                                         samples_sub_directory='JPEGImages',
                                                         targets_sub_directory='SegmentationClass', augment=True,
                                                         dataset_hyper_params=dataset_params, cache_labels=cache_labels,
                                                         cache_images=cache_images)

        self.valset = PascalVOC2012SegmentationDataSet(root=self.root_dir,
                                                       list_file='ImageSets/Segmentation/val.txt',
                                                       samples_sub_directory='JPEGImages',
                                                       targets_sub_directory='SegmentationClass', augment=True,
                                                       dataset_hyper_params=dataset_params, cache_labels=cache_labels,
                                                       cache_images=cache_images)
        self.classes = self.trainset.classes


class PascalAUG2012SegmentationDataSetInterface(DatasetInterface):
    def __init__(self, dataset_params=None, cache_labels=False, cache_images=False):
        if dataset_params is None:
            dataset_params = dict()
        super().__init__(dataset_params=dataset_params)

        self.root_dir = dataset_params['dataset_dir'] if 'dataset_dir' in dataset_params.keys() \
            else '/data/pascal_voc_2012/VOCaug/dataset/'

        self.trainset = PascalAUG2012SegmentationDataSet(
            root=self.root_dir,
            list_file='trainval.txt',
            samples_sub_directory='img',
            targets_sub_directory='cls', augment=True,
            dataset_hyper_params=dataset_params, cache_labels=cache_labels,
            cache_images=cache_images)

        self.valset = PascalAUG2012SegmentationDataSet(
            root=self.root_dir,
            list_file='val.txt',
            samples_sub_directory='img',
            targets_sub_directory='cls', augment=False,
            dataset_hyper_params=dataset_params, cache_labels=cache_labels,
            cache_images=cache_images)

        self.classes = self.trainset.classes


class CoCoDataSetInterfaceBase(DatasetInterface):
    def __init__(self, dataset_params=None):
        if dataset_params is None:
            dataset_params = dict()
        super().__init__(dataset_params=dataset_params)

        self.root_dir = dataset_params['dataset_dir'] if 'dataset_dir' in dataset_params.keys() else '/data/coco/'


class CoCoSegmentationDatasetInterface(CoCoDataSetInterfaceBase):
    def __init__(self, dataset_params=None, cache_labels: bool = False, cache_images: bool = False,
                 dataset_classes_inclusion_tuples_list: list = None):
        super().__init__(dataset_params=dataset_params)

        # backwards compatability patch for legacy dataset params
        img_size = core_utils.get_param(dataset_params, "img_size")
        crop_size = core_utils.get_param(dataset_params, "crop_size")

        train_transforms = [RandomFlip(),
                            Rescale(long_size=img_size),
                            RandomRescale(scales=(0.5, 2.0)),
                            PadShortToCropSize(crop_size=crop_size),
                            CropImageAndMask(crop_size=crop_size, mode="random")]
        val_transforms = [Rescale(short_size=crop_size),
                          CropImageAndMask(crop_size=crop_size, mode="center")]

        self.trainset = CoCoSegmentationDataSet(
            root_dir=self.root_dir,
            list_file='instances_train2017.json',
            samples_sub_directory='images/train2017',
            targets_sub_directory='annotations',
            cache_labels=cache_labels,
            cache_images=cache_images,
            transforms=train_transforms,
            dataset_classes_inclusion_tuples_list=dataset_classes_inclusion_tuples_list)

        self.valset = CoCoSegmentationDataSet(
            root_dir=self.root_dir,
            list_file='instances_val2017.json',
            samples_sub_directory='images/val2017',
            targets_sub_directory='annotations',
            cache_labels=cache_labels,
            cache_images=cache_images,
            transforms=val_transforms,
            dataset_classes_inclusion_tuples_list=dataset_classes_inclusion_tuples_list)

        self.coco_classes = self.trainset.classes


class CityscapesDatasetInterface(DatasetInterface):
    def __init__(self, dataset_params=None, cache_labels: bool = False, cache_images: bool = False):
        super().__init__(dataset_params=dataset_params)
        root_dir = core_utils.get_param(dataset_params, "dataset_dir", "/data/cityscapes")
        image_mask_transforms = core_utils.get_param(dataset_params, "image_mask_transforms")
        image_mask_transforms_aug = core_utils.get_param(dataset_params, "image_mask_transforms_aug")

        # Backwards compatability fix for SegmentationDataset refactor
        train_transforms = image_mask_transforms_aug['Compose']['transforms']
        val_transforms = image_mask_transforms['Compose']['transforms']

        self.trainset = CityscapesDataset(
            root_dir=root_dir,
            list_file='lists/train.lst',
            labels_csv_path="lists/labels.csv",
            cache_labels=cache_labels,
            cache_images=cache_images,
            transforms=train_transforms)

        self.valset = CityscapesDataset(
            root_dir=root_dir,
            list_file='lists/val.lst',
            labels_csv_path="lists/labels.csv",
            cache_labels=cache_labels,
            cache_images=cache_images,
            transforms=val_transforms)

        self.classes = self.trainset.classes


class SuperviselyPersonsDatasetInterface(DatasetInterface):
    def __init__(self, dataset_params=None, cache_labels: bool = False, cache_images: bool = False):
        super().__init__(dataset_params=dataset_params)
        root_dir = get_param(dataset_params, "dataset_dir", "/data/supervisely-persons")

        self.trainset = SuperviselyPersonsDataset(
            root_dir=root_dir,
            list_file='train.csv',
            dataset_hyper_params=dataset_params,
            cache_labels=cache_labels,
            cache_images=cache_images,
            image_mask_transforms_aug=get_param(dataset_params, "image_mask_transforms_aug", transforms.Compose([])),
            augment=True
        )

        self.valset = SuperviselyPersonsDataset(
            root_dir=root_dir,
            list_file='val.csv',
            dataset_hyper_params=dataset_params,
            cache_labels=cache_labels,
            cache_images=cache_images,
            image_mask_transforms=get_param(dataset_params, "image_mask_transforms", transforms.Compose([])),
            augment=False
        )

        self.classes = self.trainset.classes


class DetectionDatasetInterface(DatasetInterface):
    def build_data_loaders(self, batch_size_factor=1, num_workers=8, train_batch_size=None, val_batch_size=None,
                           test_batch_size=None, distributed_sampler: bool = False):

        train_sampler = InfiniteSampler(self.trainset, seed=0)

        train_batch_sampler = BatchSampler(
            sampler=train_sampler,
            batch_size=self.dataset_params.batch_size,
            drop_last=False,
        )

        self.train_loader = DataLoader(self.trainset,
                                       batch_sampler=train_batch_sampler,
                                       num_workers=num_workers,
                                       pin_memory=True,
                                       worker_init_fn=worker_init_reset_seed,
                                       collate_fn=self.dataset_params.train_collate_fn)

        if distributed_sampler:
            sampler = torch.utils.data.distributed.DistributedSampler(self.valset, shuffle=False)
        else:
            sampler = torch.utils.data.SequentialSampler(self.valset)

        val_loader = torch.utils.data.DataLoader(self.valset,
                                                 num_workers=num_workers,
                                                 pin_memory=True,
                                                 sampler=sampler,
                                                 batch_size=self.dataset_params.val_batch_size,
                                                 collate_fn=self.dataset_params.val_collate_fn)

        self.val_loader = val_loader


class PascalVOCUnifiedDetectionDatasetInterface(DetectionDatasetInterface):

    def __init__(self, dataset_params=None):
        if dataset_params is None:
            dataset_params = dict()
        super().__init__(dataset_params=dataset_params)

        self.data_dir = self.dataset_params.data_dir
        train_input_dim = (self.dataset_params.train_image_size, self.dataset_params.train_image_size)
        val_input_dim = (self.dataset_params.val_image_size, self.dataset_params.val_image_size)
        train_max_num_samples = get_param(self.dataset_params, "train_max_num_samples")
        val_max_num_samples = get_param(self.dataset_params, "val_max_num_samples")

        if self.dataset_params.download:
            PascalVOCDetectionDataset.download(data_dir=self.data_dir)

        train_dataset_names = ["train2007", "val2007", "train2012", "val2012"]
        # We divide train_max_num_samples between the datasets
        if train_max_num_samples:
            max_num_samples_per_train_dataset = [len(segment) for segment in
                                                 np.array_split(range(train_max_num_samples), len(train_dataset_names))]
        else:
            max_num_samples_per_train_dataset = [None] * len(train_dataset_names)
        train_sets = [PascalVOCDetectionDataset(data_dir=self.data_dir,
                                                input_dim=train_input_dim,
                                                cache=self.dataset_params.cache_train_images,
                                                cache_dir=self.dataset_params.cache_dir,
                                                transforms=self.dataset_params.train_transforms,
                                                images_sub_directory='images/' + trainset_name + '/',
                                                class_inclusion_list=self.dataset_params.class_inclusion_list,
                                                max_num_samples=max_num_samples_per_train_dataset[i])
                      for i, trainset_name in enumerate(train_dataset_names)]

        testset2007 = PascalVOCDetectionDataset(data_dir=self.data_dir,
                                                input_dim=val_input_dim,
                                                cache=self.dataset_params.cache_val_images,
                                                cache_dir=self.dataset_params.cache_dir,
                                                transforms=self.dataset_params.val_transforms,
                                                images_sub_directory='images/test2007/',
                                                class_inclusion_list=self.dataset_params.class_inclusion_list,
                                                max_num_samples=val_max_num_samples)

        self.classes = train_sets[1].classes
        self.trainset = ConcatDataset(train_sets)
        self.valset = testset2007

        self.trainset.collate_fn = self.dataset_params.train_collate_fn
        self.trainset.classes = self.classes
        self.trainset.img_size = self.dataset_params.train_image_size
        self.trainset.cache_labels = self.dataset_params.cache_train_images


class CoCoDetectionDatasetInterface(DetectionDatasetInterface):
    def __init__(self, dataset_params={}):
        super(CoCoDetectionDatasetInterface, self).__init__(dataset_params=dataset_params)

        # IF CACHE- CREATING THE CACHE FILE WILL HAPPEN ONLY FOR RANK 0, THEN ALL THE OTHER RANKS SIMPLY READ FROM IT.
        local_rank = get_local_rank()
        with wait_for_the_master(local_rank):
            self.trainset = COCODetectionDataset(data_dir=self.dataset_params.data_dir,
                                                 subdir=self.dataset_params.train_subdir,
                                                 json_file=self.dataset_params.train_json_file,
                                                 input_dim=self.dataset_params.train_input_dim,
                                                 cache=self.dataset_params.cache_train_images,
                                                 cache_dir=self.dataset_params.cache_dir,
                                                 transforms=self.dataset_params.train_transforms,
                                                 tight_box_rotation=self.dataset_params.tight_box_rotation,
                                                 class_inclusion_list=self.dataset_params.class_inclusion_list,
                                                 max_num_samples=self.dataset_params.train_max_num_samples,
                                                 with_crowd=False)

        # IF CACHE- CREATING THE CACHE FILE WILL HAPPEN ONLY FOR RANK 0, THEN ALL THE OTHER RANKS SIMPLY READ FROM IT.
        with wait_for_the_master(local_rank):
            self.valset = COCODetectionDataset(
                data_dir=self.dataset_params.data_dir,
                json_file=self.dataset_params.val_json_file,
                subdir=self.dataset_params.val_subdir,
                cache_dir=self.dataset_params.cache_dir,
                cache=self.dataset_params.cache_val_images,
                input_dim=self.dataset_params.val_input_dim,
                transforms=self.dataset_params.val_transforms,
                class_inclusion_list=self.dataset_params.class_inclusion_list,
                max_num_samples=self.dataset_params.val_max_num_samples,
                with_crowd=self.dataset_params.with_crowd)
        self.classes = COCO_DETECTION_CLASSES_LIST

@@ -4,6 +4,7 @@ from omegaconf import DictConfig
 from torch.utils.data import DataLoader
 
 from super_gradients.common import MultiGPUMode
+from super_gradients.training.dataloaders import dataloaders
 from super_gradients.training.models import SgModule
 from super_gradients.training.models.all_architectures import KD_ARCHITECTURES
 from super_gradients.training.models.kd_modules.kd_module import KDModule
@@ -57,8 +58,14 @@ class KDTrainer(Trainer):
 
         trainer = KDTrainer(**kwargs)
 
-        # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-        trainer.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+        # INSTANTIATE DATA LOADERS
+        train_dataloader = dataloaders.get(name=cfg.train_dataloader,
+                                           dataset_params=cfg.dataset_params.train_dataset_params,
+                                           dataloader_params=cfg.dataset_params.train_dataloader_params)
+
+        val_dataloader = dataloaders.get(name=cfg.val_dataloader,
+                                         dataset_params=cfg.dataset_params.val_dataset_params,
+                                         dataloader_params=cfg.dataset_params.val_dataloader_params)
 
         student = models.get(cfg.student_architecture, arch_params=cfg.student_arch_params,
                              strict_load=cfg.student_checkpoint_params.strict_load,
@@ -75,7 +82,8 @@ class KDTrainer(Trainer):
         # TRAIN
         trainer.train(training_params=cfg.training_hyperparams, student=student, teacher=teacher,
                       kd_architecture=cfg.architecture, kd_arch_params=cfg.arch_params,
-                      run_teacher_on_eval=cfg.run_teacher_on_eval)
+                      run_teacher_on_eval=cfg.run_teacher_on_eval,
+                      train_loader=train_dataloader, valid_loader=val_dataloader)
 
     def build_model(self,
                     # noqa: C901 - too complex
@@ -303,7 +311,8 @@ class KDTrainer(Trainer):
                                    })
         return hyper_param_config
 
-    def _instantiate_ema_model(self, decay: float = 0.9999, beta: float = 15, exp_activation: bool = True) -> KDModelEMA:
+    def _instantiate_ema_model(self, decay: float = 0.9999, beta: float = 15,
+                               exp_activation: bool = True) -> KDModelEMA:
         """Instantiate KD ema model for KDModule.
 
         If the model is of class KDModule, the instance will be adapted to work on knowledge distillation.
@@ -330,7 +339,8 @@ class KDTrainer(Trainer):
 
     def train(self, model: KDModule = None, training_params: dict = dict(), student: SgModule = None,
               teacher: torch.nn.Module = None, kd_architecture: Union[KDModule.__class__, str] = 'kd_module',
-              kd_arch_params: dict = dict(), run_teacher_on_eval=False, *args, **kwargs):
+              kd_arch_params: dict = dict(), run_teacher_on_eval=False, train_loader: DataLoader = None,
+              valid_loader: DataLoader = None, *args, **kwargs):
         """
         Trains the student network (wrapped in KDModule network).
 
@@ -342,6 +352,8 @@ class KDTrainer(Trainer):
         :param kd_architecture: KDModule architecture to use, currently only 'kd_module' is supported (default='kd_module').
         :param kd_arch_params: architecture params to pas to kd_architecture constructor.
         :param run_teacher_on_eval: bool- whether to run self.teacher at eval mode regardless of self.train(mode)
+        :param train_loader: Dataloader for train set.
+        :param valid_loader: Dataloader for validation.
         """
         kd_net = self.net or model
         if kd_net is None:
@@ -352,4 +364,5 @@ class KDTrainer(Trainer):
                                               run_teacher_on_eval=run_teacher_on_eval,
                                               student=student,
                                               teacher=teacher)
-        super(KDTrainer, self).train(model=kd_net, training_params=training_params)
+        super(KDTrainer, self).train(model=kd_net, training_params=training_params,
+                                     train_loader=train_loader, valid_loader=valid_loader)

@@ -1,5 +1,5 @@
 from super_gradients.training.utils import HpmStruct
-
+from copy import deepcopy
 DEFAULT_TRAINING_PARAMS = {"lr_warmup_epochs": 0,
                            "lr_cooldown_epochs": 0,
                            "warmup_initial_lr": None,
@@ -99,7 +99,8 @@ class TrainingParams(HpmStruct):
 
     def __init__(self, **entries):
         # WE initialize by the default training params, overridden by the provided params
-        super().__init__(**DEFAULT_TRAINING_PARAMS)
+        default_training_params = deepcopy(DEFAULT_TRAINING_PARAMS)
+        super().__init__(**default_training_params)
         self.set_schema(TRAINING_PARAM_SCHEMA)
         if len(entries) > 0:
             self.override(**entries)

@@ -23,14 +23,13 @@ from super_gradients.training.models.all_architectures import ARCHITECTURES
 from super_gradients.common.decorators.factory_decorator import resolve_param
 from super_gradients.common.environment import env_helpers
 from super_gradients.common.abstractions.abstract_logger import get_logger
-from super_gradients.common.factories.datasets_factory import DatasetsFactory
 from super_gradients.common.factories.list_factory import ListFactory
 from super_gradients.common.factories.losses_factory import LossesFactory
 from super_gradients.common.factories.metrics_factory import MetricsFactory
 from super_gradients.common.sg_loggers import SG_LOGGERS
 from super_gradients.common.sg_loggers.abstract_sg_logger import AbstractSGLogger
 from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
-from super_gradients.training import utils as core_utils, models
+from super_gradients.training import utils as core_utils, models, dataloaders
 from super_gradients.training.models import SgModule
 from super_gradients.training.pretrained_models import PRETRAINED_NUM_CLASSES
 from super_gradients.training.utils import sg_trainer_utils
@@ -38,7 +37,6 @@ from super_gradients.training.utils.quantization_utils import QATCallback
 from super_gradients.training.utils.sg_trainer_utils import MonitoredValue, parse_args
 from super_gradients.training.exceptions.sg_trainer_exceptions import UnsupportedOptimizerFormat, \
     IllegalDataloaderInitialization
-from super_gradients.training.datasets import DatasetInterface
 from super_gradients.training.losses import LOSSES
 from super_gradients.training.metrics.metric_utils import get_metrics_titles, get_metrics_results_tuple, \
     get_logging_values, \
@@ -203,8 +201,14 @@ class Trainer:
 
         trainer = Trainer(**kwargs)
 
-        # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-        trainer.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+        # INSTANTIATE DATA LOADERS
+        train_dataloader = dataloaders.get(name=cfg.train_dataloader,
+                                           dataset_params=cfg.dataset_params.train_dataset_params,
+                                           dataloader_params=cfg.dataset_params.train_dataloader_params)
+
+        val_dataloader = dataloaders.get(name=cfg.val_dataloader,
+                                         dataset_params=cfg.dataset_params.val_dataset_params,
+                                         dataloader_params=cfg.dataset_params.val_dataloader_params)
 
         # BUILD NETWORK
         model = models.get(name=cfg.architecture,
@@ -217,7 +221,10 @@ class Trainer:
                            )
 
         # TRAIN
-        trainer.train(model=model, training_params=cfg.training_hyperparams)
+        trainer.train(model=model,
+                      train_loader=train_dataloader,
+                      valid_loader=val_dataloader,
+                      training_params=cfg.training_hyperparams)
 
     def _set_dataset_properties(self, classes, test_loader, train_loader, valid_loader):
         if any([train_loader, valid_loader, classes]) and not all([train_loader, valid_loader, classes]):
@@ -239,23 +246,6 @@ class Trainer:
         self.dataset_params, self.train_loader, self.valid_loader, self.test_loader, self.classes = \
             HpmStruct(**dataset_params), train_loader, valid_loader, test_loader, classes
 
-    @resolve_param('dataset_interface', DatasetsFactory())
-    def connect_dataset_interface(self, dataset_interface: DatasetInterface, data_loader_num_workers: int = 8):
-        """
-        :param dataset_interface: DatasetInterface object
-        :param data_loader_num_workers: The number of threads to initialize the Data Loaders with
-            The dataset to be connected
-        """
-        if self.train_loader:
-            logger.warning("Overriding the dataloaders that Trainer was initialized with")
-        self.dataset_interface = dataset_interface
-        self.train_loader, self.valid_loader, self.test_loader, self.classes = \
-            self.dataset_interface.get_data_loaders(batch_size_factor=self.num_devices,
-                                                    num_workers=data_loader_num_workers,
-                                                    distributed_sampler=self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)
-
-        self.dataset_params = self.dataset_interface.get_dataset_params()
-
     # FIXME - we need to resolve flake8's 'function is too complex' for this function
     @deprecated(target=None, deprecated_in='2.3.0', remove_in='3.0.0')
     def build_model(self,  # noqa: C901 - too complex
@@ -533,11 +523,7 @@ class Trainer:
 
     def _prep_net_for_train(self):
         if self.arch_params is None:
-            default_arch_params = HpmStruct(sync_bn=False)
-            arch_params = getattr(self.net, "arch_params", default_arch_params)
-            self.arch_params = default_arch_params
-            if arch_params is not None:
-                self.arch_params.override(**arch_params.to_dict())
+            self._init_arch_params()
 
         # TODO: REMOVE THE BELOW LINE (FOR BACKWARD COMPATIBILITY)
         if self.checkpoint_params is None:
@@ -555,8 +541,16 @@ class Trainer:
         self.external_checkpoint_path = core_utils.get_param(self.training_params, "resume_path")
         self._load_checkpoint_to_model()
 
+    def _init_arch_params(self):
+        default_arch_params = HpmStruct(sync_bn=False)
+        arch_params = getattr(self.net, "arch_params", default_arch_params)
+        self.arch_params = default_arch_params
+        if arch_params is not None:
+            self.arch_params.override(**arch_params.to_dict())
+
     # FIXME - we need to resolve flake8's 'function is too complex' for this function
-    def train(self, model: nn.Module = None, training_params: dict = dict(), train_loader: DataLoader = None, valid_loader: DataLoader = None):  # noqa: C901
+    def train(self, model: nn.Module = None, training_params: dict = None, train_loader: DataLoader = None,
+              valid_loader: DataLoader = None):  # noqa: C901
         """
 
         train - Trains the Model
@@ -835,6 +829,8 @@ class Trainer:
         :return:
         """
         global logger
+        if training_params is None:
+            training_params = dict()
 
         self.train_loader = train_loader or self.train_loader
         self.valid_loader = valid_loader or self.valid_loader
@@ -1427,13 +1423,16 @@ class Trainer:
                              "calling test or through training_params when calling train(...)")
         if self.test_loader is None:
             raise ValueError("Test dataloader is required to perform test. Make sure to either pass it through "
-                             "test_loader arg or calling connect_dataset_interface upon a DatasetInterface instance "
-                             "with a non empty testset attribute.")
+                             "test_loader arg.")
 
         # RESET METRIC RUNNERS
         self._reset_metrics()
         self.test_metrics.to(self.device)
 
+        if self.arch_params is None:
+            self._init_arch_params()
+        self._net_to_device()
+
     def _add_metrics_update_callback(self, phase: Phase):
         """
         Adds MetricsUpdateCallback to be fired at phase

@@ -5,8 +5,7 @@ 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 (
+from super_gradients.training.dataloaders.dataloaders import (
     cifar10_train,
     cifar10_val,
     cifar100_train,
@@ -15,24 +14,6 @@ from super_gradients.training.dataloaders.dataloader_factory import (
 
 
 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")
@@ -52,24 +33,6 @@ class TestCifarTrainer(unittest.TestCase):
             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")

import super_gradients
import torch
import unittest
import numpy as np
from PIL import Image
import tensorflow.keras as keras
from super_gradients.training import MultiGPUMode, models
from super_gradients.training import Trainer
from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ExternalDatasetInterface, \
    ImageNetDatasetInterface
from super_gradients.training.metrics import Accuracy, Top5


class DataGenerator(keras.utils.Sequence):
    def __init__(self, samples, batch_size=1, dims=(320, 320), n_channels=3,
                 n_classes=1000, shuffle=True):
        self.dims = dims
        self.batch_size = batch_size
        self.samples = samples
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        # Fraction of dataset to be used - for faster testing
        fraction_of_dataset = 0.01
        return int(np.floor(len(self.samples) / self.batch_size) * fraction_of_dataset)

    def __getitem__(self, index):
        indices = self.indices[index * self.batch_size:(index + 1) * self.batch_size]
        list_IDs_temp = [self.samples[k] for k in indices]
        X, y = self.__data_generation(list_IDs_temp)
        return X, y

    def on_epoch_end(self):
        self.indices = np.arange(len(self.samples))
        if self.shuffle:
            np.random.shuffle(self.indices)

    def __data_generation(self, list_IDs_temp):
        X = np.empty((self.batch_size, *self.dims, self.n_channels), dtype=np.float32)
        y = np.empty((self.batch_size), dtype=int)
        for i, ID in enumerate(list_IDs_temp):
            image = Image.open(ID[0])
            image = image.resize((self.dims))
            rgb_image = Image.new("RGB", image.size)
            rgb_image.paste(image)
            X[i, ] = np.array(rgb_image)
            y[i] = ID[1]
        return X, keras.utils.to_categorical(y, num_classes=self.n_classes)


def create_imagenet_dataset():
    dataset_params = {"batch_size": 1}
    dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params=dataset_params)
    return dataset


class TransposeCollateFn(object):

    def __init__(self, new_shape):
        self.new_shape = new_shape

    def __call__(self, batch):
        new_inputs = []
        new_targets = []
        for img in batch:
            squeezed_input = img[0].squeeze(axis=0)
            transposed_data = np.transpose(squeezed_input, self.new_shape)
            new_inputs.append(torch.from_numpy(transposed_data))
            argmax_target = np.argmax(img[1], 1)
            new_targets.append(torch.from_numpy(argmax_target))
        return torch.stack(new_inputs, 0), torch.cat(new_targets, 0)


class TestExternalDatasetInterface(unittest.TestCase):

    def setUp(self):
        super_gradients.init_trainer()
        dataset = create_imagenet_dataset()
        data_samples_train = dataset.trainset.samples
        data_samples_val = dataset.valset.samples
        # batch size: 1 is only for the creation of the external keras loader
        self.keras_params = {'dims': (256, 256),
                             'batch_size': 1,
                             'n_classes': 1000,
                             'n_channels': 3,
                             'shuffle': True}
        training_generator = DataGenerator(data_samples_train, **self.keras_params)
        testing_generator = DataGenerator(data_samples_val, **self.keras_params)
        external_num_classes = 1000
        collate_fn = TransposeCollateFn((2, 0, 1))
        self.external_dataset_params = {'batch_size': 16,
                                        'test_batch_size': 16,
                                        'train_collate_fn': collate_fn,
                                        'val_collate_fn': collate_fn}
        self.test_external_dataset_interface = ExternalDatasetInterface(train_loader=training_generator,
                                                                        val_loader=testing_generator,
                                                                        num_classes=external_num_classes,
                                                                        dataset_params=self.external_dataset_params)

    def test_transpose_collate_fn(self):
        collate_fn = TransposeCollateFn((2, 0, 1))
        dims = self.keras_params['dims']
        n_channels = self.keras_params['n_channels']
        batch_size = self.external_dataset_params['batch_size']
        dummy_batch = []
        dummy_input = np.expand_dims(np.random.rand(dims[0], dims[1], n_channels), axis=0)
        dummy_target = np.expand_dims(np.random.rand(1), axis=0)
        for i in range(batch_size):
            dummy_batch.append((dummy_input, dummy_target))
        collate_fn_output = collate_fn.__call__(dummy_batch)
        dummy_tensor = torch.rand(batch_size, n_channels, dims[0], dims[1])
        self.assertEqual(dummy_tensor.shape, collate_fn_output[0].shape)

    def test_model_train(self):
        train_params = {"max_epochs": 2, "lr_decay_factor": 0.1, "initial_lr": 0.025,
                        "loss": "cross_entropy",
                        "train_metrics_list": [Accuracy(), Top5()],
                        "valid_metrics_list": [Accuracy(), Top5()],
                        "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                        "greater_metric_to_watch_is_better": True}

        arch_params = {'num_classes': 1000}
        trainer = Trainer("test", model_checkpoints_location='local',
                          multi_gpu=MultiGPUMode.OFF)
        trainer.connect_dataset_interface(dataset_interface=self.test_external_dataset_interface,
                                          data_loader_num_workers=8)
        model = models.get("resnet50", arch_params)
        trainer.train(model=model, training_params=train_params)


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

@@ -6,7 +6,8 @@ from super_gradients.training import models
 import super_gradients
 import torch
 import os
-from super_gradients import Trainer, ClassificationTestDatasetInterface
+from super_gradients import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
 
 
@@ -38,32 +39,32 @@ class TestTrainer(unittest.TestCase):
     @staticmethod
     def get_classification_trainer(name=''):
         trainer = Trainer(name, model_checkpoints_location='local')
-        dataset_params = {"batch_size": 4}
-        dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params, image_size=224)
-        trainer.connect_dataset_interface(dataset)
-        model = models.get("resnet18", arch_params={"num_classes": 5})
+        model = models.get("resnet18", num_classes=5)
         return trainer, model
 
     def test_train(self):
         trainer, model = self.get_classification_trainer(self.folder_names[0])
-        trainer.train(model=model, training_params=self.training_params)
+        trainer.train(model=model, training_params=self.training_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
     def test_save_load(self):
         trainer, model = self.get_classification_trainer(self.folder_names[1])
-        trainer.train(model=model, training_params=self.training_params)
-
+        trainer.train(model=model, training_params=self.training_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         resume_training_params = self.training_params.copy()
         resume_training_params["resume"] = True
         resume_training_params["max_epochs"] = 2
         trainer, model = self.get_classification_trainer(self.folder_names[1])
-        trainer.train(model=model, training_params=resume_training_params)
+        trainer.train(model=model, training_params=resume_training_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
     def test_checkpoint_content(self):
         """VERIFY THAT ALL CHECKPOINTS ARE SAVED AND CONTAIN ALL THE EXPECTED KEYS"""
         trainer, model = self.get_classification_trainer(self.folder_names[5])
         params = self.training_params.copy()
         params["save_ckpt_epoch_list"] = [1]
-        trainer.train(model=model, training_params=params)
+        trainer.train(model=model, training_params=params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         ckpt_filename = ['ckpt_best.pth', 'ckpt_latest.pth', 'ckpt_epoch_1.pth']
         ckpt_paths = [os.path.join(trainer.checkpoints_dir_path, suf) for suf in ckpt_filename]
         for ckpt_path in ckpt_paths:

@@ -1,7 +1,6 @@
 # PACKAGE IMPORTS FOR EXTERNAL USAGE
 
-from tests.integration_tests.s3_dataset_test import TestDataset
 from tests.integration_tests.ema_train_integration_test import EMAIntegrationTest
 from tests.integration_tests.lr_test import LRTest
 
-_all__ = [TestDataset, EMAIntegrationTest, LRTest]
+__all__ = ["EMAIntegrationTest", "LRTest"]

@@ -4,11 +4,9 @@ import re
 
 from super_gradients.training import models
 
-from super_gradients import (
-    Trainer,
-    ClassificationTestDatasetInterface,
-    SegmentationTestDatasetInterface,
-)
+from super_gradients import Trainer
+from super_gradients.training.dataloaders.dataloaders import segmentation_test_dataloader, \
+    classification_test_dataloader
 from super_gradients.training.utils.callbacks import ModelConversionCheckCallback
 from super_gradients.training.metrics import Accuracy, Top5, IoU
 from super_gradients.training.losses.stdc_loss import STDCLoss
@@ -16,7 +14,6 @@ from super_gradients.training.losses.ddrnet_loss import DDRNetLoss
 
 from deci_lab_client.models import ModelMetadata, HardwareType, FrameworkType
 
-
 checkpoint_dir = "/Users/daniel/Documents/LALA"
 
 
@@ -44,6 +41,8 @@ def generate_model_metadata(architecture: str, task: Task):
 
 CLASSIFICATION = ["efficientnet_b0", "regnetY200", "regnetY400", "regnetY600", "regnetY800", "mobilenet_v3_large"]
 SEMANTIC_SEGMENTATION = ["ddrnet_23", "stdc1_seg", "stdc2_seg", "regseg48"]
+
+
 # TODO: ADD YOLOX ARCHITECTURES AND TESTS
 
 
@@ -70,13 +69,12 @@ class ConversionCallbackTest(unittest.TestCase):
                 "phase_callbacks": phase_callbacks,
             }
 
-            trainer = Trainer(f"{architecture}_example", model_checkpoints_location="local", ckpt_root_dir=checkpoint_dir)
-            dataset = ClassificationTestDatasetInterface(dataset_params={"batch_size": 10})
-
-            trainer.connect_dataset_interface(dataset, data_loader_num_workers=0)
+            trainer = Trainer(f"{architecture}_example", model_checkpoints_location="local",
+                              ckpt_root_dir=checkpoint_dir)
             model = models.get(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
             try:
-                trainer.train(model=model, training_params=train_params)
+                trainer.train(model=model, training_params=train_params, train_loader=classification_test_dataloader(),
+                              valid_loader=classification_test_dataloader())
             except Exception as e:
                 self.fail(f"Model training didn't succeed due to {e}")
             else:
@@ -104,10 +102,9 @@ class ConversionCallbackTest(unittest.TestCase):
 
         for architecture in SEMANTIC_SEGMENTATION:
             model_meta_data = generate_model_metadata(architecture=architecture, task=Task.SEMANTIC_SEGMENTATION)
-            dataset = SegmentationTestDatasetInterface(dataset_params={"batch_size": 10})
-            trainer = Trainer(f"{architecture}_example", model_checkpoints_location="local", ckpt_root_dir=checkpoint_dir)
-            trainer.connect_dataset_interface(dataset, data_loader_num_workers=0)
-            model = models.get(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
+            trainer = Trainer(f"{architecture}_example", model_checkpoints_location="local",
+                              ckpt_root_dir=checkpoint_dir)
+            model = models.get(name=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
 
             phase_callbacks = [
                 ModelConversionCheckCallback(model_meta_data=model_meta_data, opset_version=11, rtol=1, atol=1),
@@ -131,7 +128,8 @@ class ConversionCallbackTest(unittest.TestCase):
             train_params.update(custom_config)
 
             try:
-                trainer.train(model=model, training_params=train_params)
+                trainer.train(model=model, training_params=train_params, train_loader=segmentation_test_dataloader(image_size=512),
+                              valid_loader=segmentation_test_dataloader(image_size=512))
             except Exception as e:
                 self.fail(f"Model training didn't succeed for {architecture} due to {e}")
             else:

@@ -1,6 +1,6 @@
 import unittest
-from super_gradients import Trainer, \
-    ClassificationTestDatasetInterface
+from super_gradients import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
 from torch.optim import SGD
@@ -11,8 +11,6 @@ from deci_lab_client.models import Metric, QuantizationLevel, ModelMetadata, Opt
 class DeciLabUploadTest(unittest.TestCase):
     def setUp(self) -> None:
         self.trainer = Trainer("deci_lab_export_test_model", model_checkpoints_location='local')
-        dataset = ClassificationTestDatasetInterface(dataset_params={"batch_size": 10})
-        self.trainer.connect_dataset_interface(dataset)
 
     def test_train_with_deci_lab_integration(self):
         model_meta_data = ModelMetadata(name='model_for_deci_lab_upload_test',
@@ -49,7 +47,8 @@ class DeciLabUploadTest(unittest.TestCase):
                         "phase_callbacks": [model_conversion_callback, deci_lab_callback]}
         self.optimizer = SGD(params=net.parameters(), lr=0.1)
 
-        self.trainer.train(model=net, training_params=train_params)
+        self.trainer.train(model=net, training_params=train_params,
+                           train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader())
 
         # CLEANUP
 

@@ -1,6 +1,6 @@
-from super_gradients import ClassificationTestDatasetInterface
 from super_gradients.training import MultiGPUMode, models
 from super_gradients.training import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
 import unittest
 
@@ -25,8 +25,6 @@ class EMAIntegrationTest(unittest.TestCase):
     def _init_model(self) -> None:
         self.trainer = Trainer("resnet18_cifar_ema_test", model_checkpoints_location='local',
                                device='cpu', multi_gpu=MultiGPUMode.OFF)
-        dataset_interface = ClassificationTestDatasetInterface({"batch_size": 32})
-        self.trainer.connect_dataset_interface(dataset_interface, 8)
         self.model = models.get("resnet18_cifar", arch_params={"num_classes": 5})
 
     @classmethod
@@ -65,7 +63,9 @@ class EMAIntegrationTest(unittest.TestCase):
         self.trainer.test = CallWrapper(self.trainer.test, check_before=before_test)
         self.trainer._train_epoch = CallWrapper(self.trainer._train_epoch, check_before=before_train_epoch)
 
-        self.trainer.train(model=self.model, training_params=training_params)
+        self.trainer.train(model=self.model, training_params=training_params,
+                           train_loader=classification_test_dataloader(),
+                           valid_loader=classification_test_dataloader())
 
         self.assertIsNotNone(self.trainer.ema_model)
 

@@ -4,7 +4,8 @@ import os
 
 from super_gradients.training import models
 
-from super_gradients import Trainer, ClassificationTestDatasetInterface
+from super_gradients import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
 
 
@@ -30,10 +31,7 @@ class LRTest(unittest.TestCase):
     @staticmethod
     def get_trainer(name=''):
         trainer = Trainer(name, model_checkpoints_location='local')
-        dataset_params = {"batch_size": 4}
-        dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        trainer.connect_dataset_interface(dataset)
-        model = models.get("resnet18_cifar", arch_params={"num_classes": 5})
+        model = models.get("resnet18_cifar", num_classes=5)
         return trainer, model
 
     def test_function_lr(self):
@@ -44,22 +42,25 @@ class LRTest(unittest.TestCase):
 
         # test if we are able that lr_function supports functions with this structure
         training_params = {**self.training_params, "lr_mode": "function", "lr_schedule_function": test_lr_function}
-        trainer.train(model=model, training_params=training_params)
-
+        trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         # test that we assert lr_function is callable
         training_params = {**self.training_params, "lr_mode": "function"}
         with self.assertRaises(AssertionError):
-            trainer.train(model=model, training_params=training_params)
+            trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
+                          valid_loader=classification_test_dataloader())
 
     def test_cosine_lr(self):
         trainer, model = self.get_trainer(self.folder_name)
         training_params = {**self.training_params, "lr_mode": "cosine", "cosine_final_lr_ratio": 0.01}
-        trainer.train(model=model, training_params=training_params)
+        trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
     def test_step_lr(self):
         trainer, model = self.get_trainer(self.folder_name)
         training_params = {**self.training_params, "lr_mode": "step", "lr_decay_factor": 0.1, "lr_updates": [4]}
-        trainer.train(model=model, training_params=training_params)
+        trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
 
 if __name__ == '__main__':

@@ -1,11 +1,13 @@
 import unittest
-import super_gradients
+
 from super_gradients.training import MultiGPUMode
 from super_gradients.training import Trainer
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface, \
-    ClassificationTestDatasetInterface, CityscapesDatasetInterface, SegmentationTestDatasetInterface, \
-    CoCoSegmentationDatasetInterface, DetectionTestDatasetInterface
-from super_gradients.training.utils.segmentation_utils import coco_sub_classes_inclusion_tuples_list
+from super_gradients.training.dataloaders import imagenet_val, imagenet_vit_base_val
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader, coco2017_val_yolox, \
+    coco2017_val_ssd_lite_mobilenet_v2, detection_test_dataloader, coco_segmentation_val, cityscapes_val, \
+    cityscapes_stdc_seg50_val, cityscapes_stdc_seg75_val, segmentation_test_dataloader
+from super_gradients.training.utils.detection_utils import CrowdDetectionCollateFN
+
 from super_gradients.training.metrics import Accuracy, IoU
 import os
 import shutil
@@ -13,13 +15,10 @@ from super_gradients.training.utils.ssd_utils import SSDPostPredictCallback
 from super_gradients.training.models.detection_models.ssd import DEFAULT_SSD_LITE_MOBILENET_V2_ARCH_PARAMS
 from super_gradients.training.losses.ddrnet_loss import DDRNetLoss
 from super_gradients.training.metrics import DetectionMetrics
-from super_gradients.training.transforms.transforms import Rescale
 from super_gradients.training.losses.stdc_loss import STDCLoss
 from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
-from super_gradients.training.utils.detection_utils import DetectionCollateFN, CrowdDetectionCollateFN
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import CoCoDetectionDatasetInterface
-from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
 from super_gradients.training import models
+import super_gradients
 
 
 class PretrainedModelsTest(unittest.TestCase):
@@ -34,6 +33,15 @@ class PretrainedModelsTest(unittest.TestCase):
                                                     {"image_size": (224, 224),
                                                      "patch_size": (16, 16)}}
 
+        self.imagenet_pretrained_trainsfer_learning_arch_params = {"resnet": {},
+                                                                   "regnet": {},
+                                                                   "repvgg_a0": {"build_residual_branches": True},
+                                                                   "efficientnet_b0": {},
+                                                                   "mobilenet": {},
+                                                                   "vit_base":
+                                                                       {"image_size": (224, 224),
+                                                                        "patch_size": (16, 16)}}
+
         self.imagenet_pretrained_ckpt_params = {"pretrained_weights": "imagenet"}
 
         self.imagenet21k_pretrained_ckpt_params = {"pretrained_weights": "imagenet21k"}
@@ -54,16 +62,11 @@ class PretrainedModelsTest(unittest.TestCase):
                                                "vit_large": 0.8564,
                                                "beit_base_patch16_224": 0.85
                                                }
-        self.imagenet_dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params={"batch_size": 128})
+        self.imagenet_dataset = imagenet_val(dataloader_params={"batch_size": 128})
 
-        self.imagenet_dataset_05_mean_std = ImageNetDatasetInterface(data_dir="/data/Imagenet",
-                                                                     dataset_params={"batch_size": 128,
-                                                                                     "img_mean": [0.5, 0.5, 0.5],
-                                                                                     "img_std": [0.5, 0.5, 0.5],
-                                                                                     "resize_size": 248
-                                                                                     })
+        self.imagenet_dataset_05_mean_std = imagenet_vit_base_val(dataloader_params={"batch_size": 128})
 
-        self.transfer_classification_dataset = ClassificationTestDatasetInterface(image_size=224)
+        self.transfer_classification_dataloader = classification_test_dataloader(image_size=224)
 
         self.transfer_classification_train_params = {"max_epochs": 3,
                                                      "lr_updates": [1],
@@ -83,76 +86,13 @@ class PretrainedModelsTest(unittest.TestCase):
                                             'coco_ssd_mobilenet_v1': {'num_classes': 80}}
         self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"}
 
-        from super_gradients.training.transforms.transforms import DetectionMosaic, DetectionMixup, \
-            DetectionRandomAffine, \
-            DetectionTargetsFormatTransform, DetectionPaddedRescale, DetectionHSV, DetectionHorizontalFlip
-
-        yolox_train_transforms = [DetectionMosaic(input_dim=(640, 640), prob=1.0),
-                                  DetectionRandomAffine(degrees=10., translate=0.1, scales=[0.1, 2], shear=2.0,
-                                                        target_size=(640, 640),
-                                                        filter_box_candidates=False, wh_thr=0, area_thr=0, ar_thr=0),
-                                  DetectionMixup(input_dim=(640, 640), mixup_scale=[0.5, 1.5], prob=1.0, flip_prob=0.5),
-                                  DetectionHSV(prob=1.0, hgain=5, sgain=30, vgain=30),
-                                  DetectionHorizontalFlip(prob=0.5),
-                                  DetectionPaddedRescale(input_dim=(640, 640), max_targets=120),
-                                  DetectionTargetsFormatTransform(output_format=DetectionTargetsFormat.LABEL_CXCYWH)]
-        yolox_val_transforms = [DetectionPaddedRescale(input_dim=(640, 640)),
-                                DetectionTargetsFormatTransform(max_targets=50,
-                                                                output_format=DetectionTargetsFormat.LABEL_CXCYWH)]
-
-        ssd_train_transforms = [DetectionMosaic(input_dim=(640, 640), prob=1.0),
-                                DetectionRandomAffine(degrees=0., translate=0.1, scales=[0.5, 1.5], shear=.0,
-                                                      target_size=(640, 640),
-                                                      filter_box_candidates=True, wh_thr=2, area_thr=0.1, ar_thr=20),
-                                DetectionMixup(input_dim=(640, 640), mixup_scale=[0.5, 1.5], prob=0., flip_prob=0.),
-                                DetectionHSV(prob=.0, hgain=5, sgain=30, vgain=30),
-                                DetectionHorizontalFlip(prob=0.),
-                                DetectionPaddedRescale(input_dim=(640, 640), max_targets=120),
-                                DetectionTargetsFormatTransform(output_format=DetectionTargetsFormat.LABEL_CXCYWH)]
-        ssd_val_transforms = [DetectionPaddedRescale(input_dim=(640, 640)),
-                              DetectionTargetsFormatTransform(max_targets=50,
-                                                              output_format=DetectionTargetsFormat.LABEL_CXCYWH)]
-
         self.coco_dataset = {
-            'yolox': CoCoDetectionDatasetInterface(
-                dataset_params={"data_dir": "/data/coco",
-                                "train_subdir": "images/train2017",
-                                "val_subdir": "images/val2017",
-                                "train_json_file": "instances_train2017.json",
-                                "val_json_file": "instances_val2017.json",
-                                "batch_size": 16,
-                                "val_batch_size": 128,
-                                "val_image_size": 640,
-                                "train_image_size": 640,
-                                "train_transforms": yolox_train_transforms,
-                                "val_transforms": yolox_val_transforms,
-
-                                "val_collate_fn": CrowdDetectionCollateFN(),
-                                "train_collate_fn": DetectionCollateFN(),
-                                "cache_dir_path": None,
-                                "cache_train_images": False,
-                                "cache_val_images": False,
-                                "with_crowd": True}),
-
-            'ssd_mobilenet': CoCoDetectionDatasetInterface(dataset_params={"data_dir": "/data/coco",
-                                                                           "train_subdir": "images/train2017",
-                                                                           "val_subdir": "images/val2017",
-                                                                           "train_json_file": "instances_train2017.json",
-                                                                           "val_json_file": "instances_val2017.json",
-                                                                           "batch_size": 16,
-                                                                           "val_batch_size": 128,
-                                                                           "val_image_size": 320,
-                                                                           "train_image_size": 320,
-                                                                           "train_transforms": ssd_train_transforms,
-                                                                           "val_transforms": ssd_val_transforms,
-
-                                                                           "val_collate_fn": CrowdDetectionCollateFN(),
-                                                                           "train_collate_fn": DetectionCollateFN(),
-                                                                           "cache_dir_path": None,
-                                                                           "cache_train_images": False,
-                                                                           "cache_val_images": False,
-                                                                           "with_crowd": True})
-        }
+            'yolox': coco2017_val_yolox(dataloader_params={"collate_fn": CrowdDetectionCollateFN()},
+                                        dataset_params={"with_crowd": True}),
+
+            'ssd_mobilenet': coco2017_val_ssd_lite_mobilenet_v2(
+                dataloader_params={"collate_fn": CrowdDetectionCollateFN()},
+                dataset_params={"with_crowd": True})}
 
         self.coco_pretrained_maps = {'ssd_lite_mobilenet_v2': 0.2052,
                                      'coco_ssd_mobilenet_v1': 0.243,
@@ -162,70 +102,60 @@ class PretrainedModelsTest(unittest.TestCase):
                                      "yolox_n": 0.2677,
                                      "yolox_t": 0.3718}
 
-        self.transfer_detection_dataset = DetectionTestDatasetInterface(image_size=320, classes=['class1', 'class2'])
+        self.transfer_detection_dataset = detection_test_dataloader()
 
         ssd_dboxes = DEFAULT_SSD_LITE_MOBILENET_V2_ARCH_PARAMS['anchors']
-        self.transfer_detection_train_params = {
-            'ssd_lite_mobilenet_v2':
-                {
-                    "max_epochs": 3,
-                    "lr_mode": "cosine",
-                    "initial_lr": 0.01,
-                    "cosine_final_lr_ratio": 0.01,
-                    "lr_warmup_epochs": 3,
-                    "batch_accumulate": 1,
-                    "loss": "ssd_loss",
-                    "criterion_params": {"dboxes": ssd_dboxes},
-                    "optimizer": "SGD",
-                    "warmup_momentum": 0.8,
-                    "optimizer_params": {"momentum": 0.937,
-                                         "weight_decay": 0.0005,
-                                         "nesterov": True},
-                    "train_metrics_list": [],
-                    "valid_metrics_list": [
-                        DetectionMetrics(
-                            post_prediction_callback=SSDPostPredictCallback(),
-                            num_cls=len(self.transfer_detection_dataset.classes))],
-                    "loss_logging_items_names": ['smooth_l1', 'closs', 'Loss'],
-                    "metric_to_watch": "mAP@0.50:0.95",
-                    "greater_metric_to_watch_is_better": True
-                },
-            "yolox":
-                {"max_epochs": 3,
-                 "lr_mode": "cosine",
-                 "cosine_final_lr_ratio": 0.05,
-                 "warmup_bias_lr": 0.0,
-                 "warmup_momentum": 0.9,
-                 "initial_lr": 0.02,
-                 "loss": "yolox_loss",
-                 "criterion_params": {
-                     "strides": [8, 16, 32],  # output strides of all yolo outputs
-                     "num_classes": len(self.transfer_detection_dataset.classes)},
-
-                 "loss_logging_items_names": ["iou", "obj", "cls", "l1", "num_fg", "Loss"],
-
-                 "train_metrics_list": [],
-                 "valid_metrics_list": [
-                     DetectionMetrics(
-                         post_prediction_callback=YoloPostPredictionCallback(),
-                         normalize_targets=True,
-                         num_cls=len(self.transfer_detection_dataset.classes))],
-                 "metric_to_watch": 'mAP@0.50:0.95',
-                 "greater_metric_to_watch_is_better": True}
+        self.transfer_detection_train_params_ssd = {
+            "max_epochs": 3,
+            "lr_mode": "cosine",
+            "initial_lr": 0.01,
+            "cosine_final_lr_ratio": 0.01,
+            "lr_warmup_epochs": 3,
+            "batch_accumulate": 1,
+            "loss": "ssd_loss",
+            "criterion_params": {"dboxes": ssd_dboxes},
+            "optimizer": "SGD",
+            "warmup_momentum": 0.8,
+            "optimizer_params": {"momentum": 0.937,
+                                 "weight_decay": 0.0005,
+                                 "nesterov": True},
+            "train_metrics_list": [],
+            "valid_metrics_list": [
+                DetectionMetrics(
+                    post_prediction_callback=SSDPostPredictCallback(),
+                    num_cls=5)],
+            "loss_logging_items_names": ['smooth_l1', 'closs', 'Loss'],
+            "metric_to_watch": "mAP@0.50:0.95",
+            "greater_metric_to_watch_is_better": True
         }
+        self.transfer_detection_train_params_yolox = {"max_epochs": 3,
+                                                      "lr_mode": "cosine",
+                                                      "cosine_final_lr_ratio": 0.05,
+                                                      "warmup_bias_lr": 0.0,
+                                                      "warmup_momentum": 0.9,
+                                                      "initial_lr": 0.02,
+                                                      "loss": "yolox_loss",
+                                                      "criterion_params": {
+                                                          "strides": [8, 16, 32],  # output strides of all yolo outputs
+                                                          "num_classes": 5},
+
+                                                      "loss_logging_items_names": ["iou", "obj", "cls", "l1", "num_fg",
+                                                                                   "Loss"],
+
+                                                      "train_metrics_list": [],
+                                                      "valid_metrics_list": [
+                                                          DetectionMetrics(
+                                                              post_prediction_callback=YoloPostPredictionCallback(),
+                                                              normalize_targets=True,
+                                                              num_cls=5)],
+                                                      "metric_to_watch": 'mAP@0.50:0.95',
+                                                      "greater_metric_to_watch_is_better": True}
 
         self.coco_segmentation_subclass_pretrained_arch_params = {
             "shelfnet34_lw": {"num_classes": 21, "image_size": 512}}
         self.coco_segmentation_subclass_pretrained_ckpt_params = {"pretrained_weights": "coco_segmentation_subclass"}
         self.coco_segmentation_subclass_pretrained_mious = {"shelfnet34_lw": 0.651}
-        self.coco_segmentation_dataset = CoCoSegmentationDatasetInterface(dataset_params={
-            "batch_size": 24,
-            "val_batch_size": 24,
-            "dataset_dir": "/data/coco/",
-            "img_size": 608,
-            "crop_size": 512
-        }, dataset_classes_inclusion_tuples_list=coco_sub_classes_inclusion_tuples_list()
-        )
+        self.coco_segmentation_dataset = coco_segmentation_val()
 
         self.cityscapes_pretrained_models = ["ddrnet_23", "ddrnet_23_slim", "stdc1_seg50", "regseg48"]
         self.cityscapes_pretrained_arch_params = {
@@ -248,31 +178,12 @@ class PretrainedModelsTest(unittest.TestCase):
                                             "pp_lite_b_seg50": 0.7648,
                                             "pp_lite_b_seg75": 0.7852}
 
-        self.cityscapes_dataset = CityscapesDatasetInterface(dataset_params={
-            "batch_size": 3,
-            "val_batch_size": 3,
-            "dataset_dir": "/data/cityscapes/",
-            "crop_size": 1024,
-            "img_size": 1024,
-            "image_mask_transforms_aug": [],
-            "image_mask_transforms": []  # no transform for evaluation
-        }, cache_labels=False)
-
-        self.cityscapes_dataset_rescaled50 = CityscapesDatasetInterface(dataset_params={
-            "batch_size": 3,
-            "val_batch_size": 3,
-            "image_mask_transforms_aug": [],
-            "image_mask_transforms": [Rescale(scale_factor=0.5)]  # no transform for evaluation
-        }, cache_labels=False)
-
-        self.cityscapes_dataset_rescaled75 = CityscapesDatasetInterface(dataset_params={
-            "batch_size": 3,
-            "val_batch_size": 3,
-            "image_mask_transforms_aug": [],
-            "image_mask_transforms": [Rescale(scale_factor=0.75)]  # no transform for evaluation
-        }, cache_labels=False)
-
-        self.transfer_segmentation_dataset = SegmentationTestDatasetInterface(image_size=1024)
+        self.cityscapes_dataset = cityscapes_val()
+
+        self.cityscapes_dataset_rescaled50 = cityscapes_stdc_seg50_val()
+        self.cityscapes_dataset_rescaled75 = cityscapes_stdc_seg75_val()
+
+        self.transfer_segmentation_dataset = segmentation_test_dataloader(image_size=1024)
         self.ddrnet_transfer_segmentation_train_params = {"max_epochs": 3,
                                                           "initial_lr": 1e-2,
                                                           "loss": DDRNetLoss(),
@@ -327,6 +238,7 @@ class PretrainedModelsTest(unittest.TestCase):
             "train_metrics_list": [IoU(5)],
             "valid_metrics_list": [IoU(5)],
             "loss_logging_items_names": ["loss"],
+
             "metric_to_watch": "IoU",
             "greater_metric_to_watch_is_better": True
         }
@@ -334,154 +246,160 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_resnet50_imagenet(self):
         trainer = Trainer('imagenet_pretrained_resnet50', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         model = models.get("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet50"], delta=0.001)
 
     def test_transfer_learning_resnet50_imagenet(self):
         trainer = Trainer('imagenet_pretrained_resnet50_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_resnet34_imagenet(self):
         trainer = Trainer('imagenet_pretrained_resnet34', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet34"], delta=0.001)
 
     def test_transfer_learning_resnet34_imagenet(self):
         trainer = Trainer('imagenet_pretrained_resnet34_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_resnet18_imagenet(self):
         trainer = Trainer('imagenet_pretrained_resnet18', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet18"], delta=0.001)
 
     def test_transfer_learning_resnet18_imagenet(self):
         trainer = Trainer('imagenet_pretrained_resnet18_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_regnetY800_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY800', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY800"], delta=0.001)
 
     def test_transfer_learning_regnetY800_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY800_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_regnetY600_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY600', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY600"], delta=0.001)
 
     def test_transfer_learning_regnetY600_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY600_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_regnetY400_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY400', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY400"], delta=0.001)
 
     def test_transfer_learning_regnetY400_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY400_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_regnetY200_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY200', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY200"], delta=0.001)
 
     def test_transfer_learning_regnetY200_imagenet(self):
         trainer = Trainer('imagenet_pretrained_regnetY200_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_repvgg_a0_imagenet(self):
         trainer = Trainer('imagenet_pretrained_repvgg_a0', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["repvgg_a0"], delta=0.001)
 
     def test_transfer_learning_repvgg_a0_imagenet(self):
         trainer = Trainer('imagenet_pretrained_repvgg_a0_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         model = models.get("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_regseg48_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_regseg48', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset, data_loader_num_workers=8)
         model = models.get("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["regseg48"], delta=0.001)
@@ -489,18 +407,18 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_transfer_learning_regseg48_cityscapes(self):
         trainer = Trainer('regseg48_cityscapes_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"],
                            **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.regseg_transfer_segmentation_train_params)
+        trainer.train(model=model, train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset,
+                      training_params=self.regseg_transfer_segmentation_train_params)
 
     def test_pretrained_ddrnet23_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_ddrnet23', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset, data_loader_num_workers=8)
         model = models.get("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23"], delta=0.001)
@@ -508,10 +426,9 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_ddrnet23_slim_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_ddrnet23_slim', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset, data_loader_num_workers=8)
         model = models.get("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23_slim"], delta=0.001)
@@ -519,94 +436,94 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_transfer_learning_ddrnet23_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_ddrnet23_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
                            **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.ddrnet_transfer_segmentation_train_params)
+        trainer.train(model=model, training_params=self.ddrnet_transfer_segmentation_train_params,
+                      train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset)
 
     def test_transfer_learning_ddrnet23_slim_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_ddrnet23_slim_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
                            **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.ddrnet_transfer_segmentation_train_params)
+        trainer.train(model=model, training_params=self.ddrnet_transfer_segmentation_train_params,
+                      train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset)
 
     def test_pretrained_coco_segmentation_subclass_pretrained_shelfnet34_lw(self):
         trainer = Trainer('coco_segmentation_subclass_pretrained_shelfnet34_lw', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("shelfnet34_lw",
                            arch_params=self.coco_segmentation_subclass_pretrained_arch_params["shelfnet34_lw"],
                            **self.coco_segmentation_subclass_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.coco_segmentation_dataset.val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_segmentation_dataset,
                            test_metrics_list=[IoU(21)], metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.coco_segmentation_subclass_pretrained_mious["shelfnet34_lw"], delta=0.001)
 
     def test_pretrained_efficientnet_b0_imagenet(self):
         trainer = Trainer('imagenet_pretrained_efficientnet_b0', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["efficientnet_b0"], delta=0.001)
 
     def test_transfer_learning_efficientnet_b0_imagenet(self):
         trainer = Trainer('imagenet_pretrained_efficientnet_b0_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
+
         model = models.get("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_ssd_lite_mobilenet_v2_coco(self):
         trainer = Trainer('coco_ssd_lite_mobilenet_v2', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['ssd_mobilenet'], data_loader_num_workers=8)
         model = models.get("ssd_lite_mobilenet_v2",
                            arch_params=self.coco_pretrained_arch_params["ssd_lite_mobilenet_v2"],
                            **self.coco_pretrained_ckpt_params)
         ssd_post_prediction_callback = SSDPostPredictCallback()
-        res = trainer.test(model=model, test_loader=self.coco_dataset['ssd_mobilenet'].val_loader, test_metrics_list=[
+        res = trainer.test(model=model, test_loader=self.coco_dataset['ssd_mobilenet'], test_metrics_list=[
             DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)], metrics_progress_verbose=True)[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["ssd_lite_mobilenet_v2"], delta=0.001)
 
     def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self):
         trainer = Trainer('coco_ssd_lite_mobilenet_v2_transfer_learning',
                           model_checkpoints_location='local', multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_detection_dataset,
-                                          data_loader_num_workers=8)
         transfer_arch_params = self.coco_pretrained_arch_params['ssd_lite_mobilenet_v2'].copy()
-        transfer_arch_params['num_classes'] = len(self.transfer_detection_dataset.classes)
+        transfer_arch_params['num_classes'] = 5
         model = models.get("ssd_lite_mobilenet_v2",
                            arch_params=transfer_arch_params,
                            **self.coco_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_detection_train_params['ssd_lite_mobilenet_v2'])
+        trainer.train(model=model, training_params=self.transfer_detection_train_params_ssd,
+                      train_loader=self.transfer_detection_dataset,
+                      valid_loader=self.transfer_detection_dataset)
 
     def test_pretrained_ssd_mobilenet_v1_coco(self):
         trainer = Trainer('coco_ssd_mobilenet_v1', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['ssd_mobilenet'], data_loader_num_workers=8)
         model = models.get("ssd_mobilenet_v1",
                            arch_params=self.coco_pretrained_arch_params["coco_ssd_mobilenet_v1"],
                            **self.coco_pretrained_ckpt_params)
         ssd_post_prediction_callback = SSDPostPredictCallback()
-        res = trainer.test(model=model, test_loader=self.coco_dataset['ssd_mobilenet'].val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_dataset['ssd_mobilenet'],
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback,
-                                                               num_cls=len(
-                                                                   self.coco_dataset['ssd_mobilenet'].coco_classes))],
+                                                               num_cls=80)],
                            metrics_progress_verbose=True)[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["coco_ssd_mobilenet_v1"], delta=0.001)
 
     def test_pretrained_yolox_s_coco(self):
         trainer = Trainer('yolox_s', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
+
         model = models.get("yolox_s",
                            **self.coco_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'].val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'],
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
                                                                num_cls=80,
                                                                normalize_targets=True)])[2]
@@ -615,10 +532,9 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_yolox_m_coco(self):
         trainer = Trainer('yolox_m', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         model = models.get("yolox_m",
                            **self.coco_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'].val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'],
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
                                                                num_cls=80,
                                                                normalize_targets=True)])[2]
@@ -627,10 +543,9 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_yolox_l_coco(self):
         trainer = Trainer('yolox_l', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         model = models.get("yolox_l",
                            **self.coco_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'].val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'],
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
                                                                num_cls=80,
                                                                normalize_targets=True)])[2]
@@ -639,10 +554,10 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_yolox_n_coco(self):
         trainer = Trainer('yolox_n', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
+
         model = models.get("yolox_n",
                            **self.coco_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'].val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'],
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
                                                                num_cls=80,
                                                                normalize_targets=True)])[2]
@@ -651,10 +566,9 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_pretrained_yolox_t_coco(self):
         trainer = Trainer('yolox_t', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         model = models.get("yolox_t",
                            **self.coco_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'].val_loader,
+        res = trainer.test(model=model, test_loader=self.coco_dataset['yolox'],
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
                                                                num_cls=80,
                                                                normalize_targets=True)])[2]
@@ -664,26 +578,29 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer('test_transfer_learning_yolox_n_coco',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_detection_dataset, data_loader_num_workers=8)
-        model = models.get("yolox_n", **self.coco_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_detection_train_params["yolox"])
+        model = models.get("yolox_n", **self.coco_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_detection_train_params_yolox,
+                      train_loader=self.transfer_detection_dataset,
+                      valid_loader=self.transfer_detection_dataset)
 
     def test_transfer_learning_mobilenet_v3_large_imagenet(self):
         trainer = Trainer('imagenet_pretrained_mobilenet_v3_large_transfer_learning',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
+
         model = models.get("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_mobilenet_v3_large_imagenet(self):
         trainer = Trainer('imagenet_mobilenet_v3_large', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_large"], delta=0.001)
 
@@ -691,18 +608,20 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer('imagenet_pretrained_mobilenet_v3_small_transfer_learning',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
+
         model = models.get("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_mobilenet_v3_small_imagenet(self):
         trainer = Trainer('imagenet_mobilenet_v3_small', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_small"], delta=0.001)
 
@@ -710,28 +629,29 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = Trainer('imagenet_pretrained_mobilenet_v2_transfer_learning',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
+
         model = models.get("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
-                           **self.imagenet_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_mobilenet_v2_imagenet(self):
         trainer = Trainer('imagenet_mobilenet_v2', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
+
         model = models.get("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
                            **self.imagenet_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
+        res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v2"], delta=0.001)
 
     def test_pretrained_stdc1_seg50_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc1_seg50', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
         model = models.get("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled50.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled50,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg50"], delta=0.001)
@@ -739,18 +659,18 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_transfer_learning_stdc1_seg50_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc1_seg50_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
-                           **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params)
+                           **self.cityscapes_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params,
+                      train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset)
 
     def test_pretrained_stdc1_seg75_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc1_seg75', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
         model = models.get("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled75.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled75,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg75"], delta=0.001)
@@ -758,18 +678,18 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_transfer_learning_stdc1_seg75_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc1_seg75_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
-                           **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params)
+                           **self.cityscapes_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params,
+                      train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset)
 
     def test_pretrained_stdc2_seg50_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc2_seg50', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
         model = models.get("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled50.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled50,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg50"], delta=0.001)
@@ -777,122 +697,80 @@ class PretrainedModelsTest(unittest.TestCase):
     def test_transfer_learning_stdc2_seg50_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc2_seg50_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
-                           **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params)
+                           **self.cityscapes_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params,
+                      train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset)
 
     def test_pretrained_stdc2_seg75_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc2_seg75', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
         model = models.get("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
                            **self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled75.val_loader,
+        res = trainer.test(model=model, test_loader=self.cityscapes_dataset_rescaled75,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg75"], delta=0.001)
 
-    def test_pretrained_pplite_t_seg50_cityscapes(self):
-        trainer = Trainer('cityscapes_pretrained_pplite_t_seg50', model_checkpoints_location='local',
-                          multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
-        trainer.build_model("pp_lite_t_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
-                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled50.val_loader,
-                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                           metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg50"], delta=0.001)
-
-    def test_pretrained_pplite_t_seg75_cityscapes(self):
-        trainer = Trainer('cityscapes_pretrained_pplite_t_seg75', model_checkpoints_location='local',
-                          multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
-        trainer.build_model("pp_lite_t_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
-                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled75.val_loader,
-                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                           metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg75"], delta=0.001)
-
-    def test_pretrained_pplite_b_seg50_cityscapes(self):
-        trainer = Trainer('cityscapes_pretrained_pplite_b_seg50', model_checkpoints_location='local',
-                          multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
-        trainer.build_model("pp_lite_b_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
-                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled50.val_loader,
-                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                           metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg50"], delta=0.001)
-
-    def test_pretrained_pplite_b_seg75_cityscapes(self):
-        trainer = Trainer('cityscapes_pretrained_pplite_b_seg75', model_checkpoints_location='local',
-                          multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
-        trainer.build_model("pp_lite_b_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
-                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
-        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled75.val_loader,
-                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
-                           metrics_progress_verbose=True)[0].cpu().item()
-        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg75"], delta=0.001)
-
     def test_transfer_learning_stdc2_seg75_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc2_seg75_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         model = models.get("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
-                           **self.cityscapes_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params)
+                           **self.cityscapes_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.stdc_transfer_segmentation_train_params,
+                      train_loader=self.transfer_segmentation_dataset,
+                      valid_loader=self.transfer_segmentation_dataset)
 
     def test_transfer_learning_vit_base_imagenet21k(self):
         trainer = Trainer('imagenet21k_pretrained_vit_base',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
+
         model = models.get("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
-                           **self.imagenet21k_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet21k_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_transfer_learning_vit_large_imagenet21k(self):
         trainer = Trainer('imagenet21k_pretrained_vit_large',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
+
         model = models.get("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
-                           **self.imagenet21k_pretrained_ckpt_params)
-        trainer.train(model=model, training_params=self.transfer_classification_train_params)
+                           **self.imagenet21k_pretrained_ckpt_params, num_classes=5)
+        trainer.train(model=model, training_params=self.transfer_classification_train_params,
+                      train_loader=self.transfer_classification_dataloader,
+                      valid_loader=self.transfer_classification_dataloader)
 
     def test_pretrained_vit_base_imagenet(self):
         trainer = Trainer('imagenet_pretrained_vit_base', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset_05_mean_std, data_loader_num_workers=8)
         model = models.get("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
                            **self.imagenet_pretrained_ckpt_params)
         res = \
-            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std.val_loader,
+            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std,
                          test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_base"], delta=0.001)
 
     def test_pretrained_vit_large_imagenet(self):
         trainer = Trainer('imagenet_pretrained_vit_large', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(self.imagenet_dataset_05_mean_std, data_loader_num_workers=8)
         model = models.get("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
                            **self.imagenet_pretrained_ckpt_params)
         res = \
-            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std.val_loader,
+            trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std,

@@ -1,6 +1,6 @@
 import unittest
 
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training import Trainer, MultiGPUMode, models
 from super_gradients.training.metrics.classification_metrics import Accuracy
 import os
@@ -9,12 +9,9 @@ from super_gradients.training.utils.quantization_utils import PostQATConversionC
 
 class QATIntegrationTest(unittest.TestCase):
     def _get_trainer(self, experiment_name):
-        dataset_params = {"batch_size": 10}
-        dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
         trainer = Trainer(experiment_name,
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
-        trainer.connect_dataset_interface(dataset)
         model = models.get("resnet18", pretrained_weights="imagenet")
         return trainer, model
 
@@ -47,7 +44,8 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(model=net, training_params=train_params)
+        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
+                    valid_loader=classification_test_dataloader())
 
     def test_qat_transition(self):
         model, net = self._get_trainer("test_qat_transition")
@@ -59,7 +57,8 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(model=net, training_params=train_params)
+        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
+                    valid_loader=classification_test_dataloader())
 
     def test_qat_from_calibrated_ckpt(self):
         model, net = self._get_trainer("generate_calibrated_model")
@@ -71,7 +70,8 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(model=net, training_params=train_params)
+        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
+                    valid_loader=classification_test_dataloader())
 
         calibrated_model_path = os.path.join(model.checkpoints_dir_path, "ckpt_calibrated_percentile_99.99.pth")
 
@@ -85,7 +85,8 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(model=net, training_params=train_params)
+        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
+                    valid_loader=classification_test_dataloader())
 
 
 if __name__ == '__main__':

import unittest
import os
import shutil
from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationDatasetInterface


class TestDataset(unittest.TestCase):

    def test_donwload_dataset(self):
        default_dataset_params = {"dataset_dir": os.path.expanduser("~/test_data/"),
                                  "s3_link": "s3://research-data1/data.zip"}

        dataset = ClassificationDatasetInterface(dataset_params=default_dataset_params)

        test_sample = dataset.get_test_sample()
        self.assertListEqual([3, 64, 64], list(test_sample[0].shape))
        shutil.rmtree(default_dataset_params["dataset_dir"])


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

@@ -1,5 +1,4 @@
 # PACKAGE IMPORTS FOR EXTERNAL USAGE
-from tests.unit_tests.dataset_interface_test import TestDatasetInterface
 from tests.unit_tests.factories_test import FactoriesTest
 from tests.unit_tests.strictload_enum_test import StrictLoadEnumTest
 from tests.unit_tests.zero_weight_decay_on_bias_bn_test import ZeroWdForBnBiasTest
@@ -18,7 +17,7 @@ from tests.unit_tests.conv_bn_relu_test import TestConvBnRelu
 from tests.unit_tests.initialize_with_dataloaders_test import InitializeWithDataloadersTest
 from tests.unit_tests.training_params_factory_test import TrainingParamsTest
 
-__all__ = ['TestDatasetInterface', 'ZeroWdForBnBiasTest', 'SaveCkptListUnitTest',
+__all__ = ['ZeroWdForBnBiasTest', 'SaveCkptListUnitTest',
            'AllArchitecturesTest', 'TestAverageMeter', 'TestRepVgg', 'TestWithoutTrainTest',
            'StrictLoadEnumTest', 'TrainWithInitializedObjectsTest', 'TestAutoAugment',
            'OhemLossTest', 'EarlyStopTest', 'SegmentationTransformsTest', 'PretrainedModelsUnitTest', 'TestConvBnRelu',

@@ -4,7 +4,7 @@ import pkg_resources
 import yaml
 from torch.utils.data import DataLoader
 
-from super_gradients.training.dataloaders.dataloader_factory import cityscapes_train, cityscapes_val, \
+from super_gradients.training.dataloaders.dataloaders import cityscapes_train, cityscapes_val, \
     cityscapes_stdc_seg50_train, cityscapes_stdc_seg50_val, cityscapes_stdc_seg75_val, cityscapes_ddrnet_train, \
     cityscapes_regseg48_val, cityscapes_regseg48_train, cityscapes_ddrnet_val, cityscapes_stdc_seg75_train
 from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset

@@ -5,7 +5,7 @@ import pkg_resources
 import yaml
 from torch.utils.data import DataLoader
 
-from super_gradients.training.dataloaders.dataloader_factory import coco_segmentation_train, coco_segmentation_val
+from super_gradients.training.dataloaders.dataloaders import coco_segmentation_train, coco_segmentation_val
 from super_gradients.training.datasets.segmentation_datasets.coco_segmentation import CoCoSegmentationDataSet
 
 

@@ -2,7 +2,7 @@ import unittest
 
 from torch.utils.data import DataLoader, TensorDataset
 
-from super_gradients.training.dataloaders.dataloader_factory import (
+from super_gradients.training.dataloaders.dataloaders import (
     classification_test_dataloader,
     detection_test_dataloader,
     segmentation_test_dataloader,

import unittest

from super_gradients.training.datasets.dataset_interfaces.dataset_interface import PascalVOCUnifiedDetectionDatasetInterface,\
    CoCoDetectionDatasetInterface
from super_gradients.training.transforms.transforms import DetectionPaddedRescale, DetectionTargetsFormatTransform, DetectionMosaic, DetectionRandomAffine,\
    DetectionHSV
from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
from super_gradients.training.utils.detection_utils import DetectionCollateFN
from super_gradients.training.utils import sg_trainer_utils
from super_gradients.training import utils as core_utils


class TestDatasetInterface(unittest.TestCase):
    def setUp(self) -> None:
        self.root_dir = "/home/louis.dupont/data/"
        self.train_batch_size, self.val_batch_size = 16, 32
        self.train_image_size, self.val_image_size = 640, 640
        self.train_input_dim = (self.train_image_size, self.train_image_size)
        self.val_input_dim = (self.val_image_size, self.val_image_size)
        self.train_max_num_samples = 100
        self.val_max_num_samples = 90

    def setup_pascal_voc_interface(self):
        """setup PascalVOCUnifiedDetectionDatasetInterface and return dataloaders"""
        dataset_params = {
            "data_dir": self.root_dir + "pascal_unified_coco_format/",
            "cache_dir": self.root_dir + "pascal_unified_coco_format/",

            "batch_size": self.train_batch_size,
            "val_batch_size": self.val_batch_size,
            "train_image_size": self.train_image_size,
            "val_image_size": self.val_image_size,
            "train_max_num_samples": self.train_max_num_samples,
            "val_max_num_samples": self.val_max_num_samples,

            "train_transforms": [
                DetectionMosaic(input_dim=self.train_input_dim, prob=1),
                DetectionRandomAffine(degrees=0.373, translate=0.245, scales=0.898, shear=0.602, target_size=self.train_input_dim),
                DetectionHSV(prob=1, hgain=0.0138, sgain=0.664, vgain=0.464),
                DetectionPaddedRescale(input_dim=self.train_input_dim, max_targets=100),
                DetectionTargetsFormatTransform(input_format=DetectionTargetsFormat.XYXY_LABEL,
                                                output_format=DetectionTargetsFormat.LABEL_CXCYWH)],
            "val_transforms": [
                DetectionPaddedRescale(input_dim=self.val_input_dim),
                DetectionTargetsFormatTransform(input_format=DetectionTargetsFormat.XYXY_LABEL,
                                                output_format=DetectionTargetsFormat.LABEL_CXCYWH)],
            "train_collate_fn": DetectionCollateFN(),
            "val_collate_fn": DetectionCollateFN(),
            "download": False,
            "cache_train_images": False,
            "cache_val_images": False,
            "class_inclusion_list": ["person"]
        }
        dataset_interface = PascalVOCUnifiedDetectionDatasetInterface(dataset_params=dataset_params)
        train_loader, valid_loader, _test_loader, _classes = dataset_interface.get_data_loaders()
        return train_loader, valid_loader

    def setup_coco_detection_interface(self):
        """setup CoCoDetectionDatasetInterface and return dataloaders"""
        dataset_params = {
            "data_dir": "/data/coco",
            "train_subdir": "images/train2017",  # sub directory path of data_dir containing the train data.
            "val_subdir": "images/val2017",  # sub directory path of data_dir containing the validation data.
            "train_json_file": "instances_train2017.json",  # path to coco train json file, data_dir/annotations/train_json_file.
            "val_json_file": "instances_val2017.json",  # path to coco validation json file, data_dir/annotations/val_json_file.

            "batch_size": self.train_batch_size,
            "val_batch_size": self.val_batch_size,
            "train_image_size": self.train_image_size,
            "val_image_size": self.val_image_size,
            "train_max_num_samples": self.train_max_num_samples,
            "val_max_num_samples": self.val_max_num_samples,

            "mixup_prob": 1.0,  # probability to apply per-sample mixup
            "degrees": 10.,  # rotation degrees, randomly sampled from [-degrees, degrees]
            "shear": 2.0,  # shear degrees, randomly sampled from [-degrees, degrees]
            "flip_prob": 0.5,  # probability to apply horizontal flip
            "hsv_prob": 1.0,  # probability to apply HSV transform
            "hgain": 5,  # HSV transform hue gain (randomly sampled from [-hgain, hgain])
            "sgain": 30,  # HSV transform saturation gain (randomly sampled from [-sgain, sgain])
            "vgain": 30,  # HSV transform value gain (randomly sampled from [-vgain, vgain])
            "mosaic_scale": [0.1, 2],  # random rescale range (keeps size by padding/cropping) after mosaic transform.
            "mixup_scale": [0.5, 1.5],  # random rescale range for the additional sample in mixup
            "mosaic_prob": 1.,  # probability to apply mosaic
            "translate": 0.1,  # image translation fraction
            "filter_box_candidates": False,  # whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.
            "wh_thr": 2,  # edge size threshold when filter_box_candidates = True (pixels)
            "ar_thr": 20,  # aspect ratio threshold when filter_box_candidates = True
            "area_thr": 0.1,  # threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True
            "tight_box_rotation": False,

            "download": False,
            "train_collate_fn": DetectionCollateFN(),
            "val_collate_fn": DetectionCollateFN(),
            "cache_train_images": False,
            "cache_val_images": False,
            "cache_dir": "/home/data/cache",  # Depends on the user
            "class_inclusion_list": None
            # "with_crowd": True
        }

        dataset_interface = CoCoDetectionDatasetInterface(dataset_params=dataset_params)
        train_loader, valid_loader, _test_loader, _classes = dataset_interface.get_data_loaders()
        return train_loader, valid_loader

    def test_coco_detection(self):
        """Check that the dataset interface is correctly instantiated, and that the batch items are of expected size"""
        train_loader, valid_loader = self.setup_coco_detection_interface()

        for loader, batch_size, image_size, max_num_samples in [(train_loader, self.train_batch_size, self.train_image_size, self.train_max_num_samples),
                                                                (valid_loader, self.val_batch_size, self.val_image_size, self.val_max_num_samples)]:

            # The dataset is at most of length max_num_samples, but can be smaller if not enough samples
            self.assertGreaterEqual(max_num_samples, len(loader.dataset))

            batch_items = next(iter(loader))
            batch_items = core_utils.tensor_container_to_device(batch_items, 'cuda', non_blocking=True)

            inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)
            self.assertListEqual([batch_size, 3, image_size, image_size], list(inputs.shape))

    def test_pascal_voc(self):
        """Check that the dataset interface is correctly instantiated, and that the batch items are of expected size"""
        train_loader, valid_loader = self.setup_pascal_voc_interface()

        for loader, batch_size, image_size, max_num_samples in [(train_loader, self.train_batch_size, self.train_image_size, self.train_max_num_samples),
                                                                (valid_loader, self.val_batch_size, self.val_image_size, self.val_max_num_samples)]:
            # The dataset is at most of length max_num_samples, but can be smaller if not enough samples
            self.assertGreaterEqual(max_num_samples, len(loader.dataset))

            batch_items = next(iter(loader))
            batch_items = core_utils.tensor_container_to_device(batch_items, 'cuda', non_blocking=True)

            inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)
            self.assertListEqual([batch_size, 3, image_size, image_size], list(inputs.shape))


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

@@ -1,12 +1,10 @@
 import unittest
 
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import CoCoDetectionDatasetInterface
+from super_gradients.training.dataloaders.dataloaders import coco2017_train, coco2017_val
 from super_gradients.training.metrics.detection_metrics import DetectionMetrics
 
 from super_gradients.training import Trainer, models
 from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
-from super_gradients.training.utils.detection_utils import CrowdDetectionCollateFN, DetectionCollateFN, \
-    DetectionTargetsFormat
 
 
 class TestDatasetStatisticsTensorboardLogger(unittest.TestCase):
@@ -19,44 +17,11 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase):
         browser and make sure the text and plots in the tensorboard are as expected.
         """
         # Create dataset
-        dataset = CoCoDetectionDatasetInterface(dataset_params={"data_dir": "/data/coco",
-                                                                "train_subdir": "images/train2017",
-                                                                "val_subdir": "images/val2017",
-                                                                "train_json_file": "instances_train2017.json",
-                                                                "val_json_file": "instances_val2017.json",
-                                                                "batch_size": 16,
-                                                                "val_batch_size": 128,
-                                                                "val_image_size": 640,
-                                                                "train_image_size": 640,
-                                                                "hgain": 5,
-                                                                "sgain": 30,
-                                                                "vgain": 30,
-                                                                "mixup_prob": 1.0,
-                                                                "degrees": 10.,
-                                                                "shear": 2.0,
-                                                                "flip_prob": 0.5,
-                                                                "hsv_prob": 1.0,
-                                                                "mosaic_scale": [0.1, 2],
-                                                                "mixup_scale": [0.5, 1.5],
-                                                                "mosaic_prob": 1.,
-                                                                "translate": 0.1,
-                                                                "val_collate_fn": CrowdDetectionCollateFN(),
-                                                                "train_collate_fn": DetectionCollateFN(),
-                                                                "cache_dir_path": None,
-                                                                "cache_train_images": False,
-                                                                "cache_val_images": False,
-                                                                "targets_format": DetectionTargetsFormat.LABEL_CXCYWH,
-                                                                "with_crowd": True,
-                                                                "filter_box_candidates": False,
-                                                                "wh_thr": 0,
-                                                                "ar_thr": 0,
-                                                                "area_thr": 0
-                                                                })
 
         trainer = Trainer('dataset_statistics_visual_test',
                           model_checkpoints_location='local',
                           post_prediction_callback=YoloPostPredictionCallback())
-        trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
+
         model = models.get("yolox_s")
 
         training_params = {"max_epochs": 1,  # we dont really need the actual training to run
@@ -74,7 +39,7 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase):
                            "loss_logging_items_names": ["iou", "obj", "cls", "l1", "num_fg", "Loss"],
                            "metric_to_watch": "mAP@0.50:0.95",
                            }
-        trainer.train(model=model, training_params=training_params)
+        trainer.train(model=model, training_params=training_params, train_loader=coco2017_train(), valid_loader=coco2017_val())
 
 
 if __name__ == '__main__':

@@ -2,58 +2,23 @@ import os
 import unittest
 
 from super_gradients.training import Trainer, utils as core_utils, models
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import CoCoDetectionDatasetInterface
+from super_gradients.training.dataloaders.dataloaders import coco2017_val
 from super_gradients.training.datasets.datasets_conf import COCO_DETECTION_CLASSES_LIST
 from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
-from super_gradients.training.utils.detection_utils import DetectionVisualization, DetectionCollateFN, DetectionTargetsFormat
+from super_gradients.training.utils.detection_utils import DetectionVisualization
 
 
 class TestDetectionUtils(unittest.TestCase):
     def test_visualization(self):
-        # Create dataset
-        dataset = CoCoDetectionDatasetInterface(dataset_params={"data_dir": "/data/coco",
-                                                                "train_subdir": "images/train2017",
-                                                                "val_subdir": "images/val2017",
-                                                                "train_json_file": "instances_train2017.json",
-                                                                "val_json_file": "instances_val2017.json",
-                                                                "batch_size": 16,
-                                                                "val_batch_size": 4,
-                                                                "val_image_size": 640,
-                                                                "train_image_size": 640,
-                                                                "hgain": 5,
-                                                                "sgain": 30,
-                                                                "vgain": 30,
-                                                                "mixup_prob": 1.0,
-                                                                "degrees": 10.,
-                                                                "shear": 2.0,
-                                                                "flip_prob": 0.5,
-                                                                "hsv_prob": 1.0,
-                                                                "mosaic_scale": [0.1, 2],
-                                                                "mixup_scale": [0.5, 1.5],
-                                                                "mosaic_prob": 1.,
-                                                                "translate": 0.1,
-                                                                "val_collate_fn": DetectionCollateFN(),
-                                                                "train_collate_fn": DetectionCollateFN(),
-                                                                "cache_dir_path": None,
-                                                                "cache_train_images": False,
-                                                                "cache_val_images": False,
-                                                                "targets_format": DetectionTargetsFormat.LABEL_NORMALIZED_CXCYWH,
-                                                                "with_crowd": False,
-                                                                "filter_box_candidates": False,
-                                                                "wh_thr": 0,
-                                                                "ar_thr": 0,
-                                                                "area_thr": 0
-                                                                })
 
         # Create Yolo model
         trainer = Trainer('visualization_test',
                           model_checkpoints_location='local',
                           post_prediction_callback=YoloPostPredictionCallback())
-        trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
         model = models.get("yolox_n", pretrained_weights="coco")
 
         # Simulate one iteration of validation subset
-        valid_loader = trainer.valid_loader
+        valid_loader = coco2017_val()
         batch_i, (imgs, targets) = 0, next(iter(valid_loader))
         imgs = core_utils.tensor_container_to_device(imgs, trainer.device)
         targets = core_utils.tensor_container_to_device(targets, trainer.device)

@@ -2,10 +2,10 @@ import torch
 import torch.nn as nn
 import unittest
 
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.utils.early_stopping import EarlyStop
 from super_gradients.training.utils.callbacks import Phase
 from super_gradients.training.sg_trainer import Trainer
-from super_gradients.training.datasets.dataset_interfaces import ClassificationTestDatasetInterface
 from super_gradients.training.models.classification_models.resnet import ResNet18
 from super_gradients.training.metrics import Accuracy, Top5
 from torchmetrics.metric import Metric
@@ -43,8 +43,6 @@ class LossTest(nn.Module):
 class EarlyStopTest(unittest.TestCase):
     def setUp(self) -> None:
         # batch_size is equal to length of dataset, to have only one step per epoch, to ease the test.
-        dataset_params = {"batch_size": 10}
-        self.dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params, batch_size=10)
         self.net = ResNet18(num_classes=5, arch_params={})
         self.max_epochs = 10
         self.train_params = {"max_epochs": self.max_epochs, "lr_updates": [1], "lr_decay_factor": 0.1,
@@ -61,7 +59,6 @@ class EarlyStopTest(unittest.TestCase):
         epochs.
         """
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", patience=3, verbose=True)
         phase_callbacks = [early_stop_loss]
@@ -71,8 +68,8 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
-
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         excepted_end_epoch = 5
 
         # count divided by 2, because loss counter used for both train and eval.
@@ -84,8 +81,6 @@ class EarlyStopTest(unittest.TestCase):
         epochs.
         """
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
-
         early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=3,
                                    verbose=True)
         phase_callbacks = [early_stop_acc]
@@ -96,8 +91,8 @@ class EarlyStopTest(unittest.TestCase):
         train_params.update(
             {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
-
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         excepted_end_epoch = 6
 
         self.assertEqual(excepted_end_epoch, fake_metric.count)
@@ -107,7 +102,6 @@ class EarlyStopTest(unittest.TestCase):
         Test for mode=min metric, test that training stops after metric value reaches the `threshold` value.
         """
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", threshold=0.1, verbose=True)
         phase_callbacks = [early_stop_loss]
@@ -117,8 +111,8 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
-
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         excepted_end_epoch = 5
         # count divided by 2, because loss counter used for both train and eval.
         self.assertEqual(excepted_end_epoch, fake_loss.count // 2)
@@ -128,7 +122,6 @@ class EarlyStopTest(unittest.TestCase):
         Test for mode=max metric, test that training stops after metric value reaches the `threshold` value.
         """
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
 
         early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", threshold=0.94,
                                    verbose=True)
@@ -140,8 +133,8 @@ class EarlyStopTest(unittest.TestCase):
         train_params.update(
             {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
-
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         excepted_end_epoch = 7
 
         self.assertEqual(excepted_end_epoch, fake_metric.count)
@@ -152,7 +145,6 @@ class EarlyStopTest(unittest.TestCase):
         """
         # test Nan value
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", check_finite=True,
                                     verbose=True)
@@ -163,16 +155,14 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
-
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         excepted_end_epoch = 2
 
         self.assertEqual(excepted_end_epoch, fake_loss.count // 2)
 
         # test Inf value
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
-        trainer.build_model(self.net)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", patience=3, verbose=True)
         phase_callbacks = [early_stop_loss]
@@ -182,8 +172,8 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
-
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
         excepted_end_epoch = 3
         # count divided by 2, because loss counter used for both train and eval.
         self.assertEqual(excepted_end_epoch, fake_loss.count // 2)
@@ -194,7 +184,6 @@ class EarlyStopTest(unittest.TestCase):
         current_value - min_delta > best_value
         """
         trainer = Trainer("early_stop_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
 
         early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=2,
                                    min_delta=0.1, verbose=True)
@@ -206,7 +195,8 @@ class EarlyStopTest(unittest.TestCase):
         train_params.update(
             {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks})
 
-        trainer.train(model=self.net, training_params=train_params)
+        trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
         excepted_end_epoch = 5
 

import torch
import unittest
import numpy as np
import tensorflow.keras as keras
from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ExternalDatasetInterface


class DataGenerator(keras.utils.Sequence):
    def __init__(self, batch_size=1, dim=(320, 320), n_channels=3,
                 n_classes=1000, shuffle=True):
        self.dim = dim
        self.batch_size = batch_size
        self.list_IDs = np.ones(1000)
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        dataset_len = 32
        return dataset_len

    def __getitem__(self, index):
        indices = self.indices[index * self.batch_size:(index + 1) * self.batch_size]
        list_IDs_temp = [self.list_IDs[k] for k in indices]
        X, y = self.__data_generation(list_IDs_temp)
        return X.squeeze(axis=0), y.squeeze(axis=0)

    def on_epoch_end(self):
        self.indices = np.arange(len(self.list_IDs))
        if self.shuffle:
            np.random.shuffle(self.indices)

    def __data_generation(self, list_IDs_temp):
        X = np.ones((self.batch_size, self.n_channels, *self.dim), dtype=np.float32)
        y = np.ones((self.batch_size, 1), dtype=np.float32)
        return X, y


class TestExternalDatasetInterface(unittest.TestCase):

    def setUp(self):
        params = {'dim': (256, 256),
                  'batch_size': 1,
                  'n_classes': 1000,
                  'n_channels': 3,
                  'shuffle': True}
        training_generator = DataGenerator(**params)
        testing_generator = DataGenerator(**params)
        external_num_classes = 1000
        external_dataset_params = {'batch_size': 16,
                                   "val_batch_size": 16}
        self.dim = params['dim'][0]
        self.n_channels = params['n_channels']
        self.batch_size = external_dataset_params['batch_size']
        self.val_batch_size = external_dataset_params['val_batch_size']
        self.test_external_dataset_interface = ExternalDatasetInterface(train_loader=training_generator,
                                                                        val_loader=testing_generator,
                                                                        num_classes=external_num_classes,
                                                                        dataset_params=external_dataset_params)

    def test_get_data_loaders(self):
        train_loader, val_loader, _, num_classes = self.test_external_dataset_interface.get_data_loaders()
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            self.assertListEqual([self.batch_size, self.n_channels, self.dim, self.dim], list(inputs.shape))
            self.assertListEqual([self.batch_size, 1], list(targets.shape))
            self.assertEqual(torch.Tensor, type(inputs))
            self.assertEqual(torch.Tensor, type(targets))
        for batch_idx, (inputs, targets) in enumerate(val_loader):
            self.assertListEqual([self.val_batch_size, self.n_channels, self.dim, self.dim], list(inputs.shape))
            self.assertListEqual([self.val_batch_size, 1], list(targets.shape))
            self.assertEqual(torch.Tensor, type(inputs))
            self.assertEqual(torch.Tensor, type(targets))


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

@@ -2,21 +2,17 @@ import unittest
 
 import torch
 
-from super_gradients import ClassificationTestDatasetInterface, Trainer
+from super_gradients import Trainer
+from super_gradients.training import models
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
-from super_gradients.training.models import ResNet18
 
 
 class FactoriesTest(unittest.TestCase):
 
     def test_training_with_factories(self):
         trainer = Trainer("test_train_with_factories", model_checkpoints_location='local')
-        dataset_params = {"batch_size": 10}
-        dataset = {"classification_test_dataset": {"dataset_params": dataset_params}}
-        trainer.connect_dataset_interface(dataset)
-
-        net = ResNet18(num_classes=5, arch_params={})
-        trainer.build_model(net)
+        net = models.get("resnet18", num_classes=5)
         train_params = {"max_epochs": 2,
                         "lr_updates": [1],
                         "lr_decay_factor": 0.1,
@@ -32,11 +28,12 @@ class FactoriesTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
 
-        trainer.train(model=net, training_params=train_params)
+        trainer.train(model=net, training_params=train_params,
+                      train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
         self.assertIsInstance(trainer.train_metrics.Accuracy, Accuracy)
         self.assertIsInstance(trainer.valid_metrics.Top5, Top5)
-        self.assertIsInstance(trainer.dataset_interface, ClassificationTestDatasetInterface)
         self.assertIsInstance(trainer.optimizer, torch.optim.ASGD)
 
 

@@ -1,7 +1,7 @@
 import unittest
 from super_gradients.training import Trainer, models
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy
-from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.utils.callbacks import PhaseCallback, Phase, PhaseContext
 import torch
 
@@ -27,16 +27,11 @@ def test_forward_pass_prep_fn(inputs, targets, *args, **kwargs):
 
 
 class ForwardpassPrepFNTest(unittest.TestCase):
-    def setUp(self) -> None:
-        self.dataset_params = {"batch_size": 4}
-        self.dataset = ClassificationTestDatasetInterface(dataset_params=self.dataset_params)
-        self.arch_params = {'num_classes': 10}
 
     def test_resizing_with_forward_pass_prep_fn(self):
         # Define Model
         trainer = Trainer("ForwardpassPrepFNTest")
-        trainer.connect_dataset_interface(self.dataset)
-        model = models.get("resnet18", arch_params=self.arch_params)
+        model = models.get("resnet18", num_classes=5)
 
         sizes = []
         phase_callbacks = [TestInputSizesCallback(sizes)]
@@ -49,7 +44,8 @@ class ForwardpassPrepFNTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks,
                         "pre_prediction_callback": test_forward_pass_prep_fn}
-        trainer.train(model=model, training_params=train_params)
+        trainer.train(model=model, training_params=train_params, train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
         # ALTHOUGH NOT SEEN IN HERE, THE 4TH EPOCH USES LR=1, SO THIS IS THE EXPECTED LIST AS WE COLLECT
         # THE LRS AFTER THE UPDATE

@@ -2,7 +2,7 @@ import unittest
 
 from super_gradients.training import models
 
-from super_gradients import Trainer, ClassificationTestDatasetInterface
+from super_gradients import Trainer
 import torch
 from torch.utils.data import TensorDataset, DataLoader
 from super_gradients.training.metrics import Accuracy
@@ -26,22 +26,6 @@ class InitializeWithDataloadersTest(unittest.TestCase):
         label = torch.randint(0, len(self.testcase_classes), size=(test_size,))
         self.testcase_testloader = DataLoader(TensorDataset(inp, label))
 
-    def test_interface_was_not_broken(self):
-        trainer = Trainer("test_interface", model_checkpoints_location='local')
-        dataset_params = {"batch_size": 10}
-        dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        trainer.connect_dataset_interface(dataset)
-
-        model = models.get("efficientnet_b0", arch_params={"num_classes": 5})
-        train_params = {"max_epochs": 1, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step",
-                        "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": torch.nn.CrossEntropyLoss(),
-                        "optimizer": "SGD",
-                        "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
-                        "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()],
-                        "metric_to_watch": "Accuracy",
-                        "greater_metric_to_watch_is_better": True}
-        trainer.train(model=model, training_params=train_params)
-
     def test_initialization_rules(self):
         self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           train_loader=self.testcase_trainloader)
@@ -63,9 +47,7 @@ class InitializeWithDataloadersTest(unittest.TestCase):
 
     def test_train_with_dataloaders(self):
         trainer = Trainer(experiment_name="test_name", model_checkpoints_location="local")
-
-        trainer.build_model("resnet18")
-        model = models.get("resnet18", arch_params={"num_classes": 5})
+        model = models.get("resnet18", num_classes=5)
         trainer.train(model=model,
                       training_params={"max_epochs": 2,
                                        "lr_updates": [5, 6, 12],

@@ -2,10 +2,10 @@ import unittest
 
 from super_gradients.training import models
 from super_gradients.training import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.kd_trainer import KDTrainer
 import torch
 from super_gradients.training.utils.utils import check_models_have_same_weights
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.losses.kd_losses import KDLogitsLoss
 
@@ -14,8 +14,6 @@ class KDEMATest(unittest.TestCase):
     @classmethod
     def setUp(cls):
         cls.sg_trained_teacher = Trainer("sg_trained_teacher", device='cpu')
-        cls.dataset_params = {"batch_size": 5}
-        cls.dataset = ClassificationTestDatasetInterface(dataset_params=cls.dataset_params)
 
         cls.kd_train_params = {"max_epochs": 3, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step",
                                "lr_warmup_epochs": 0, "initial_lr": 0.1,
@@ -32,12 +30,13 @@ class KDEMATest(unittest.TestCase):
         """Check that the teacher EMA is a reference to the teacher net (not a copy)."""
 
         kd_model = KDTrainer("test_teacher_ema_not_duplicated", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 1000})
         teacher = models.get('resnet50', arch_params={'num_classes': 1000},
                              pretrained_weights="imagenet")
 
-        kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher)
+        kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher,
+                       train_loader=classification_test_dataloader(),
+                       valid_loader=classification_test_dataloader())
 
         self.assertTrue(kd_model.ema_model.ema.module.teacher is kd_model.net.module.teacher)
         self.assertTrue(kd_model.ema_model.ema.module.student is not kd_model.net.module.student)
@@ -48,24 +47,26 @@ class KDEMATest(unittest.TestCase):
         # Create a KD trainer and train it
         train_params = self.kd_train_params.copy()
         kd_model = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 1000})
         teacher = models.get('resnet50', arch_params={'num_classes': 1000},
                              pretrained_weights="imagenet")
 
-        kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher)
+        kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher,
+                       train_loader=classification_test_dataloader(),
+                       valid_loader=classification_test_dataloader())
         ema_model = kd_model.ema_model.ema
         net = kd_model.net
 
         # Load the trained KD trainer
         kd_model = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 1000})
         teacher = models.get('resnet50', arch_params={'num_classes': 1000},
                              pretrained_weights="imagenet")
 
         train_params["resume"] = True
-        kd_model.train(training_params=train_params, student=student, teacher=teacher)
+        kd_model.train(training_params=train_params, student=student, teacher=teacher,
+                       train_loader=classification_test_dataloader(),
+                       valid_loader=classification_test_dataloader())
         reloaded_ema_model = kd_model.ema_model.ema
         reloaded_net = kd_model.net
 
@@ -79,7 +80,8 @@ class KDEMATest(unittest.TestCase):
         self.assertTrue(not check_models_have_same_weights(reloaded_net, ema_model))
 
         # loaded student ema == loaded  student net (since load_ema_as_net = False)
-        self.assertTrue(not check_models_have_same_weights(reloaded_ema_model.module.student, reloaded_net.module.student))
+        self.assertTrue(
+            not check_models_have_same_weights(reloaded_ema_model.module.student, reloaded_net.module.student))
 
         # loaded teacher ema == loaded teacher net (teacher always loads ema)
         self.assertTrue(check_models_have_same_weights(reloaded_ema_model.module.teacher, reloaded_net.module.teacher))

@@ -1,11 +1,12 @@
 import os
 import unittest
 from copy import deepcopy
+
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.kd_trainer.kd_trainer import KDTrainer
 import torch
 
 from super_gradients.training import models
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
 from super_gradients.training.losses.kd_losses import KDLogitsLoss
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.models.classification_models.resnet import ResNet50, ResNet18
@@ -36,9 +37,6 @@ class PreTrainingEMANetCollector(PhaseCallback):
 class KDTrainerTest(unittest.TestCase):
     @classmethod
     def setUp(cls):
-        cls.dataset_params = {"batch_size": 5}
-        cls.dataset = ClassificationTestDatasetInterface(dataset_params=cls.dataset_params)
-
         cls.kd_train_params = {"max_epochs": 3, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step",
                                "lr_warmup_epochs": 0, "initial_lr": 0.1,
                                "loss": KDLogitsLoss(torch.nn.CrossEntropyLoss()),
@@ -69,9 +67,10 @@ class KDTrainerTest(unittest.TestCase):
         sg_model = KDTrainer("test_train_kd_module_external_models", device='cpu')
         teacher_model = ResNet50(arch_params={}, num_classes=5)
         student_model = ResNet18(arch_params={}, num_classes=5)
-        sg_model.connect_dataset_interface(self.dataset)
 
-        sg_model.train(training_params=self.kd_train_params, student=deepcopy(student_model), teacher=teacher_model)
+        sg_model.train(training_params=self.kd_train_params, student=deepcopy(student_model), teacher=teacher_model,
+                       train_loader=classification_test_dataloader(),
+                       valid_loader=classification_test_dataloader())
 
         # TEACHER WEIGHT'S SHOULD REMAIN THE SAME
         self.assertTrue(
@@ -83,7 +82,6 @@ class KDTrainerTest(unittest.TestCase):
 
     def test_train_model_with_input_adapter(self):
         kd_trainer = KDTrainer("train_kd_module_with_with_input_adapter", device='cpu')
-        kd_trainer.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 5})
         teacher = models.get('resnet50', arch_params={'num_classes': 5},
                              pretrained_weights="imagenet")
@@ -96,19 +94,21 @@ class KDTrainerTest(unittest.TestCase):
         kd_arch_params = {
             "teacher_input_adapter": adapter}
         kd_trainer.train(training_params=self.kd_train_params, student=student, teacher=teacher,
-                         kd_arch_params=kd_arch_params)
+                         kd_arch_params=kd_arch_params, train_loader=classification_test_dataloader(),
+                         valid_loader=classification_test_dataloader())
 
         self.assertEqual(kd_trainer.net.module.teacher_input_adapter, adapter)
 
     def test_load_ckpt_best_for_student(self):
         kd_trainer = KDTrainer("test_load_ckpt_best", device='cpu')
-        kd_trainer.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 5})
         teacher = models.get('resnet50', arch_params={'num_classes': 5},
                              pretrained_weights="imagenet")
         train_params = self.kd_train_params.copy()
         train_params["max_epochs"] = 1
-        kd_trainer.train(training_params=train_params, student=student, teacher=teacher)
+        kd_trainer.train(training_params=train_params, student=student, teacher=teacher,
+                         train_loader=classification_test_dataloader(),
+                         valid_loader=classification_test_dataloader())
         best_student_ckpt = os.path.join(kd_trainer.checkpoints_dir_path, "ckpt_best.pth")
 
         student_reloaded = models.get('resnet18', arch_params={'num_classes': 5},
@@ -119,14 +119,15 @@ class KDTrainerTest(unittest.TestCase):
 
     def test_load_ckpt_best_for_student_with_ema(self):
         kd_trainer = KDTrainer("test_load_ckpt_best", device='cpu')
-        kd_trainer.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 5})
         teacher = models.get('resnet50', arch_params={'num_classes': 5},
                              pretrained_weights="imagenet")
         train_params = self.kd_train_params.copy()
         train_params["max_epochs"] = 1
         train_params["ema"] = True
-        kd_trainer.train(training_params=train_params, student=student, teacher=teacher)
+        kd_trainer.train(training_params=train_params, student=student, teacher=teacher,
+                         train_loader=classification_test_dataloader(),
+                         valid_loader=classification_test_dataloader())
         best_student_ckpt = os.path.join(kd_trainer.checkpoints_dir_path, "ckpt_best.pth")
 
         student_reloaded = models.get('resnet18', arch_params={'num_classes': 5},
@@ -137,17 +138,17 @@ class KDTrainerTest(unittest.TestCase):
 
     def test_resume_kd_training(self):
         kd_trainer = KDTrainer("test_resume_training_start", device='cpu')
-        kd_trainer.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 5})
         teacher = models.get('resnet50', arch_params={'num_classes': 5},
                              pretrained_weights="imagenet")
         train_params = self.kd_train_params.copy()
         train_params["max_epochs"] = 1
-        kd_trainer.train(training_params=train_params, student=student, teacher=teacher)
+        kd_trainer.train(training_params=train_params, student=student, teacher=teacher,
+                         train_loader=classification_test_dataloader(),
+                         valid_loader=classification_test_dataloader())
         latest_net = deepcopy(kd_trainer.net)
 
         kd_trainer = KDTrainer("test_resume_training_start", device='cpu')
-        kd_trainer.connect_dataset_interface(self.dataset)
         student = models.get('resnet18', arch_params={'num_classes': 5})
         teacher = models.get('resnet50', arch_params={'num_classes': 5},
                              pretrained_weights="imagenet")
@@ -156,7 +157,9 @@ class KDTrainerTest(unittest.TestCase):
         train_params["resume"] = True
         collector = PreTrainingNetCollector()
         train_params["phase_callbacks"] = [collector]
-        kd_trainer.train(training_params=train_params, student=student, teacher=teacher)
+        kd_trainer.train(training_params=train_params, student=student, teacher=teacher,
+                         train_loader=classification_test_dataloader(),
+                         valid_loader=classification_test_dataloader())
 
         self.assertTrue(
             check_models_have_same_weights(collector.net, latest_net))

@@ -1,9 +1,9 @@
 import unittest
 from super_gradients.training import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.utils.callbacks import PhaseCallback, Phase, PhaseContext
 from super_gradients.training.utils.utils import check_models_have_same_weights
-from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.models import LeNet
 from copy import deepcopy
 
@@ -19,8 +19,6 @@ class PreTrainingEMANetCollector(PhaseCallback):
 
 class LoadCheckpointWithEmaTest(unittest.TestCase):
     def setUp(self) -> None:
-        self.dataset_params = {"batch_size": 4}
-        self.dataset = ClassificationTestDatasetInterface(dataset_params=self.dataset_params)
         self.train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step",
                              "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": 'SGD',
                              "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
@@ -32,22 +30,23 @@ class LoadCheckpointWithEmaTest(unittest.TestCase):
         # Define Model
         net = LeNet()
         trainer = Trainer("ema_ckpt_test", model_checkpoints_location='local')
-
-        trainer.connect_dataset_interface(self.dataset)
-
-        trainer.train(model=net, training_params=self.train_params)
+        trainer.train(model=net, training_params=self.train_params,
+                      train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
         ema_model = trainer.ema_model.ema
 
         # TRAIN FOR 1 MORE EPOCH AND COMPARE THE NET AT THE BEGINNING OF EPOCH 3 AND THE END OF EPOCH NUMBER 2
         net = LeNet()
         trainer = Trainer("ema_ckpt_test", model_checkpoints_location='local')
-        trainer.connect_dataset_interface(self.dataset)
+
         net_collector = PreTrainingEMANetCollector()
         self.train_params["resume"] = True
         self.train_params["max_epochs"] = 3
         self.train_params["phase_callbacks"] = [net_collector]
-        trainer.train(model=net, training_params=self.train_params)
+        trainer.train(model=net, training_params=self.train_params,
+                      train_loader=classification_test_dataloader(),
+                      valid_loader=classification_test_dataloader())
 
         reloaded_ema_model = net_collector.net.ema