Function build_dataset

classification/dataset/build.py:199–263

(split, config)

def build_dataset(split, config):
    if config.DATA.TRANSFORM == 'build_transform':
        transform = build_transform(split == 'train', config)
    elif config.DATA.TRANSFORM == 'build_transform_for_linear_probe':
        transform = build_transform_for_linear_probe(split == 'train', config)
    else:
        raise NotImplementedError
    print(split, transform)
    dataset = None
    nb_classes = None
    prefix = split
    if config.DATA.DATASET == 'imagenet' or config.DATA.DATASET == 'imagenet-real':
        if prefix == 'train' and not config.EVAL_MODE:
            root = os.path.join(config.DATA.DATA_PATH, 'train')
            dataset = ImageCephDataset(root, 'train',
                                       transform=transform,
                                       on_memory=config.DATA.IMG_ON_MEMORY)
        elif prefix == 'val':
            root = os.path.join(config.DATA.DATA_PATH, 'val')
            dataset = ImageCephDataset(root, 'val', transform=transform)
        nb_classes = 1000
    elif config.DATA.DATASET == 'imagenet22K':
        if prefix == 'train':
            if not config.EVAL_MODE:
                root = config.DATA.DATA_PATH
                dataset = ImageCephDataset(root, 'train',
                                           transform=transform,
                                           on_memory=config.DATA.IMG_ON_MEMORY)
            nb_classes = 21841
        elif prefix == 'val':
            root = os.path.join(config.DATA.DATA_PATH, 'val')
            dataset = ImageCephDataset(root, 'val', transform=transform)
            nb_classes = 1000
    elif config.DATA.DATASET == 'imagenetv2':
        from .imagenetv2 import ImageNetV2Dataset
        if prefix == 'train' and not config.EVAL_MODE:
            print(f'Only test split available for {config.DATA.DATASET}')
        else:
            dataset = ImageNetV2Dataset(variant='matched-frequency',
                                        transform=transform,
                                        location=config.DATA.DATA_PATH)
            nb_classes = 1000
    elif config.DATA.DATASET == 'imagenet_sketch':
        if prefix == 'train' and not config.EVAL_MODE:
            print(f'Only test split available for {config.DATA.DATASET}')
        else:
            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
            nb_classes = 1000
    elif config.DATA.DATASET == 'imagenet_a':
        if prefix == 'train' and not config.EVAL_MODE:
            print(f'Only test split available for {config.DATA.DATASET}')
        else:
            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
            nb_classes = 1000  # actual number of classes is 200
    elif config.DATA.DATASET == 'imagenet_r':
        if prefix == 'train' and not config.EVAL_MODE:
            print(f'Only test split available for {config.DATA.DATASET}')
        else: