cifar.py 5.95 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
        torch._C._log_api_usage_once(f"torchvision.datasets.{self.__class__.__name__}")
63

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

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

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

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

Philip Meier's avatar
Philip Meier committed
78
        self.data: Any = []
79
80
81
82
83
84
        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:
85
                entry = pickle.load(f, encoding='latin1')
86
                self.data.append(entry['data'])
87
                if 'labels' in entry:
88
                    self.targets.extend(entry['labels'])
89
                else:
90
                    self.targets.extend(entry['fine_labels'])
91

92
93
        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
94

95
96
        self._load_meta()

Philip Meier's avatar
Philip Meier committed
97
    def _load_meta(self) -> None:
98
99
100
101
102
        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:
103
            data = pickle.load(infile, encoding='latin1')
104
105
106
            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
107
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
108
109
110
111
112
113
114
        """
        Args:
            index (int): Index

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

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

        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
129
    def __len__(self) -> int:
130
        return len(self.data)
Soumith Chintala's avatar
Soumith Chintala committed
131

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

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

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

Soumith Chintala's avatar
Soumith Chintala committed
150
151

class CIFAR100(CIFAR10):
jvmancuso's avatar
jvmancuso committed
152
153
154
155
    """`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
156
    base_folder = 'cifar-100-python'
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
157
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
Soumith Chintala's avatar
Soumith Chintala committed
158
159
160
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
    train_list = [
161
        ['train', '16019d7e3df5f24257cddd939b257f8d'],
Soumith Chintala's avatar
Soumith Chintala committed
162
163
164
    ]

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