cifar.py 5.87 KB
Newer Older
Soumith Chintala's avatar
Soumith Chintala committed
1
2
3
4
from PIL import Image
import os
import os.path
import numpy as np
5
import pickle
6
import torch
Philip Meier's avatar
Philip Meier committed
7
from typing import Any, Callable, Optional, Tuple
Soumith Chintala's avatar
Soumith Chintala committed
8

9
from .vision import VisionDataset
10
from .utils import check_integrity, download_and_extract_archive
11

12

13
class CIFAR10(VisionDataset):
14
15
16
17
    """`CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.

    Args:
        root (string): Root directory of dataset where directory
18
            ``cifar-10-batches-py`` exists or will be saved to if download is set to True.
19
20
        train (bool, optional): If True, creates dataset from training set, otherwise
            creates from test set.
Tongzhou Wang's avatar
Tongzhou Wang committed
21
        transform (callable, optional): A function/transform that takes in an PIL image
22
23
24
25
26
27
28
29
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.

    """
Soumith Chintala's avatar
Soumith Chintala committed
30
    base_folder = 'cifar-10-batches-py'
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
31
    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
Soumith Chintala's avatar
Soumith Chintala committed
32
    filename = "cifar-10-python.tar.gz"
zhoumingjun's avatar
zhoumingjun committed
33
    tgz_md5 = 'c58f30108f718f92721af3b95e74349a'
Soumith Chintala's avatar
Soumith Chintala committed
34
    train_list = [
35
36
37
38
39
        ['data_batch_1', 'c99cafc152244af753f735de768cd75f'],
        ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'],
        ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'],
        ['data_batch_4', '634d18415352ddfa80567beed471001a'],
        ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'],
Soumith Chintala's avatar
Soumith Chintala committed
40
41
42
    ]

    test_list = [
43
        ['test_batch', '40351d587109b95175f43aff81a1287e'],
Soumith Chintala's avatar
Soumith Chintala committed
44
    ]
45
46
47
48
49
50
    meta = {
        'filename': 'batches.meta',
        'key': 'label_names',
        'md5': '5ff9c542aee3614f3951f8cda6e48888',
    }

Philip Meier's avatar
Philip Meier committed
51
52
53
54
55
56
57
58
    def __init__(
            self,
            root: str,
            train: bool = True,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            download: bool = False,
    ) -> None:
59

60
61
        super(CIFAR10, self).__init__(root, transform=transform,
                                      target_transform=target_transform)
62

63
64
        self.train = train  # training set or test set

Soumith Chintala's avatar
Soumith Chintala committed
65
66
67
68
        if download:
            self.download()

        if not self._check_integrity():
69
70
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can use download=True to download it')
71

72
        if self.train:
73
74
75
76
            downloaded_list = self.train_list
        else:
            downloaded_list = self.test_list

Philip Meier's avatar
Philip Meier committed
77
        self.data: Any = []
78
79
80
81
82
83
        self.targets = []

        # now load the picked numpy arrays
        for file_name, checksum in downloaded_list:
            file_path = os.path.join(self.root, self.base_folder, file_name)
            with open(file_path, 'rb') as f:
84
                entry = pickle.load(f, encoding='latin1')
85
                self.data.append(entry['data'])
86
                if 'labels' in entry:
87
                    self.targets.extend(entry['labels'])
88
                else:
89
                    self.targets.extend(entry['fine_labels'])
90

91
92
        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
Soumith Chintala's avatar
Soumith Chintala committed
93

94
95
        self._load_meta()

Philip Meier's avatar
Philip Meier committed
96
    def _load_meta(self) -> None:
97
98
99
100
101
        path = os.path.join(self.root, self.base_folder, self.meta['filename'])
        if not check_integrity(path, self.meta['md5']):
            raise RuntimeError('Dataset metadata file not found or corrupted.' +
                               ' You can use download=True to download it')
        with open(path, 'rb') as infile:
102
            data = pickle.load(infile, encoding='latin1')
103
104
105
            self.classes = data[self.meta['key']]
        self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)}

Philip Meier's avatar
Philip Meier committed
106
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
107
108
109
110
111
112
113
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
114
        img, target = self.data[index], self.targets[index]
115

116
117
        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
118
        img = Image.fromarray(img)
Soumith Chintala's avatar
Soumith Chintala committed
119
120
121
122
123
124
125
126
127

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

Philip Meier's avatar
Philip Meier committed
128
    def __len__(self) -> int:
129
        return len(self.data)
Soumith Chintala's avatar
Soumith Chintala committed
130

Philip Meier's avatar
Philip Meier committed
131
    def _check_integrity(self) -> bool:
Soumith Chintala's avatar
Soumith Chintala committed
132
        root = self.root
133
        for fentry in (self.train_list + self.test_list):
Soumith Chintala's avatar
Soumith Chintala committed
134
135
            filename, md5 = fentry[0], fentry[1]
            fpath = os.path.join(root, self.base_folder, filename)
soumith's avatar
soumith committed
136
            if not check_integrity(fpath, md5):
Soumith Chintala's avatar
Soumith Chintala committed
137
138
139
                return False
        return True

Philip Meier's avatar
Philip Meier committed
140
    def download(self) -> None:
Soumith Chintala's avatar
Soumith Chintala committed
141
142
143
        if self._check_integrity():
            print('Files already downloaded and verified')
            return
144
        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)
Soumith Chintala's avatar
Soumith Chintala committed
145

Philip Meier's avatar
Philip Meier committed
146
    def extra_repr(self) -> str:
147
        return "Split: {}".format("Train" if self.train is True else "Test")
148

Soumith Chintala's avatar
Soumith Chintala committed
149
150

class CIFAR100(CIFAR10):
jvmancuso's avatar
jvmancuso committed
151
152
153
154
    """`CIFAR100 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.

    This is a subclass of the `CIFAR10` Dataset.
    """
Soumith Chintala's avatar
Soumith Chintala committed
155
    base_folder = 'cifar-100-python'
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
156
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
Soumith Chintala's avatar
Soumith Chintala committed
157
158
159
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
    train_list = [
160
        ['train', '16019d7e3df5f24257cddd939b257f8d'],
Soumith Chintala's avatar
Soumith Chintala committed
161
162
163
    ]

    test_list = [
164
        ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
Soumith Chintala's avatar
Soumith Chintala committed
165
    ]
166
167
168
169
170
    meta = {
        'filename': 'meta',
        'key': 'fine_label_names',
        'md5': '7973b15100ade9c7d40fb424638fde48',
    }