feature(whl): Add ImageNet dataset.

data/prepare_mini_imagenet.py

@@ -38,7 +38,8 @@ def dict2list(dic: Dict) -> List:
 
     data_load_all = {}
     data_load_all['image_data'] = np.concatenate(
-        (train_data['image_data'], val_data['image_data'], test_data['image_data']))
+        (train_data['image_data'], val_data['image_data'], test_data['image_data'])
+    )
     data_load_all['class_dict'] = copy.deepcopy(train_data['class_dict'])
     data_load_all['class_dict'].update(val_data['class_dict'])
     data_load_all['class_dict'].update(test_data['class_dict'])

fling/dataset/__init__.py

@@ -3,4 +3,5 @@
 from .mnist import MNISTDataset
 from .tiny_imagenet import TinyImagenetDataset
 from .mini_imagenet import MiniImagenetDataset
+from .imagenet import ImagenetDataset
 from .build_dataset import get_dataset

fling/dataset/imagenet.py

@@ -0,0 +1,33 @@
+from torch.utils.data import Dataset
+from torchvision.datasets import ImageNet
+
+from fling.utils import get_data_transform
+from fling.utils.registry_utils import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register('imagenet')
+class ImagenetDataset(Dataset):
+    """
+        Implementation for Imagenet dataset. You are required to download the original data file from the website \
+        manually. For detailed information, you can refer to the doc of torchvision: \
+        https://pytorch.org/vision/stable/_modules/torchvision/datasets/imagenet.html
+    """
+    default_augmentation = dict(
+        random_resized_crop=dict(size=224, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.)),
+        horizontal_flip=dict(p=0.5),
+        Normalize=dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.24, 0.225]),
+    )
+
+    def __init__(self, cfg: dict, train: bool):
+        super(ImagenetDataset, self).__init__()
+        self.train = train
+        self.cfg = cfg
+        self.transform = get_data_transform(cfg.data.transforms, train=train)
+        split = 'train' if train else 'val'
+        self.dataset = ImageNet(root=cfg.data.data_path, split=split, transform=self.transform)
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, item: int) -> dict:
+        return {'input': self.dataset[item][0], 'class_id': self.dataset[item][1]}

fling/dataset/mini_imagenet.py

@@ -17,7 +17,7 @@ class MiniImagenetDataset(Dataset):
     _mean = [x / 255.0 for x in [120.39586422, 115.59361427, 104.54012653]]
     _std = [x / 255.0 for x in [70.68188272, 68.27635443, 72.54505529]]
     default_augmentation = dict(
-        random_resized_crop=dict(size=84, scale=(0.1, 1.0), ratio=(3. / 4., 4. / 3.)),
+        random_resized_crop=dict(size=(84, 84), scale=(0.1, 1.0), ratio=(3. / 4., 4. / 3.)),
         color_jitter=dict(brightness=0.4, contrast=0.4, saturation=0.4),
         horizontal_flip=dict(p=0.5),
         Normalize=dict(mean=_mean, std=_std),

setup.py

@@ -47,7 +47,7 @@
     python_requires=">=3.7",
     install_requires=[
         'setuptools<=66.1.1', 'yapf==0.29.0', 'torch>=1.7.0', 'torchvision', 'numpy>=1.18.0', 'easydict==1.9',
-        'tensorboard>=2.10.1', 'tqdm', 'timm', 'click', 'prettytable', 'einops'
+        'tensorboard>=2.10.1', 'tqdm', 'timm', 'click', 'prettytable', 'einops', 'scipy'
     ],
     extras_require={
         'test': [