mnist.py 20 KB
Newer Older
1
from .vision import VisionDataset
2
import warnings
Tian Qi Chen's avatar
Tian Qi Chen committed
3
4
5
from PIL import Image
import os
import os.path
6
import numpy as np
Tian Qi Chen's avatar
Tian Qi Chen committed
7
8
import torch
import codecs
9
import string
10
from typing import Any, Callable, Dict, List, Optional, Tuple
11
from urllib.error import URLError
12
13
from .utils import download_and_extract_archive, extract_archive, verify_str_arg, check_integrity
import shutil
Tian Qi Chen's avatar
Tian Qi Chen committed
14

15

16
class MNIST(VisionDataset):
17
18
19
    """`MNIST <http://yann.lecun.com/exdb/mnist/>`_ Dataset.

    Args:
20
21
        root (string): Root directory of dataset where ``MNIST/processed/training.pt``
            and  ``MNIST/processed/test.pt`` exist.
22
23
24
25
26
27
28
29
30
31
        train (bool, optional): If True, creates dataset from ``training.pt``,
            otherwise from ``test.pt``.
        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.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
    """
32

33
34
35
36
37
    mirrors = [
        'http://yann.lecun.com/exdb/mnist/',
        'https://ossci-datasets.s3.amazonaws.com/mnist/',
    ]

38
    resources = [
39
40
41
42
        ("train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
        ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
        ("t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
        ("t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
Tian Qi Chen's avatar
Tian Qi Chen committed
43
    ]
44

45
46
    training_file = 'training.pt'
    test_file = 'test.pt'
47
48
49
    classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
               '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']

50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
    @property
    def train_labels(self):
        warnings.warn("train_labels has been renamed targets")
        return self.targets

    @property
    def test_labels(self):
        warnings.warn("test_labels has been renamed targets")
        return self.targets

    @property
    def train_data(self):
        warnings.warn("train_data has been renamed data")
        return self.data

    @property
    def test_data(self):
        warnings.warn("test_data has been renamed data")
        return self.data

70
71
72
73
74
75
76
77
    def __init__(
            self,
            root: str,
            train: bool = True,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            download: bool = False,
    ) -> None:
78
79
        super(MNIST, self).__init__(root, transform=transform,
                                    target_transform=target_transform)
80
        self.train = train  # training set or test set
Tian Qi Chen's avatar
Tian Qi Chen committed
81

82
83
84
85
        if self._check_legacy_exist():
            self.data, self.targets = self._load_legacy_data()
            return

Tian Qi Chen's avatar
Tian Qi Chen committed
86
87
88
89
        if download:
            self.download()

        if not self._check_exists():
90
91
            raise RuntimeError('Dataset not found.' +
                               ' You can use download=True to download it')
Tian Qi Chen's avatar
Tian Qi Chen committed
92

93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
        self.data, self.targets = self._load_data()

    def _check_legacy_exist(self):
        processed_folder_exists = os.path.exists(self.processed_folder)
        if not processed_folder_exists:
            return False

        return all(
            check_integrity(os.path.join(self.processed_folder, file)) for file in (self.training_file, self.test_file)
        )

    def _load_legacy_data(self):
        # This is for BC only. We no longer cache the data in a custom binary, but simply read from the raw data
        # directly.
        data_file = self.training_file if self.train else self.test_file
        return torch.load(os.path.join(self.processed_folder, data_file))

    def _load_data(self):
        image_file = f"{'train' if self.train else 't10k'}-images-idx3-ubyte"
        data = read_image_file(os.path.join(self.raw_folder, image_file))

        label_file = f"{'train' if self.train else 't10k'}-labels-idx1-ubyte"
        targets = read_label_file(os.path.join(self.raw_folder, label_file))

        return data, targets
Tian Qi Chen's avatar
Tian Qi Chen committed
118

119
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
120
121
122
123
124
125
126
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
127
        img, target = self.data[index], int(self.targets[index])
Tian Qi Chen's avatar
Tian Qi Chen committed
128
129
130
131
132
133
134
135
136
137
138
139
140

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.numpy(), mode='L')

        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

141
    def __len__(self) -> int:
142
        return len(self.data)
Tian Qi Chen's avatar
Tian Qi Chen committed
143

144
    @property
145
    def raw_folder(self) -> str:
146
147
148
        return os.path.join(self.root, self.__class__.__name__, 'raw')

    @property
149
    def processed_folder(self) -> str:
150
151
152
        return os.path.join(self.root, self.__class__.__name__, 'processed')

    @property
153
    def class_to_idx(self) -> Dict[str, int]:
154
155
        return {_class: i for i, _class in enumerate(self.classes)}

156
    def _check_exists(self) -> bool:
157
158
159
160
        return all(
            check_integrity(os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0]))
            for url, _ in self.resources
        )
161

162
    def download(self) -> None:
163
        """Download the MNIST data if it doesn't exist already."""
Tian Qi Chen's avatar
Tian Qi Chen committed
164
165
166
167

        if self._check_exists():
            return

168
        os.makedirs(self.raw_folder, exist_ok=True)
Tian Qi Chen's avatar
Tian Qi Chen committed
169

170
        # download files
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        for filename, md5 in self.resources:
            for mirror in self.mirrors:
                url = "{}{}".format(mirror, filename)
                try:
                    print("Downloading {}".format(url))
                    download_and_extract_archive(
                        url, download_root=self.raw_folder,
                        filename=filename,
                        md5=md5
                    )
                except URLError as error:
                    print(
                        "Failed to download (trying next):\n{}".format(error)
                    )
                    continue
                finally:
                    print()
                break
            else:
                raise RuntimeError("Error downloading {}".format(filename))
Tian Qi Chen's avatar
Tian Qi Chen committed
191

192
    def extra_repr(self) -> str:
193
        return "Split: {}".format("Train" if self.train is True else "Test")
194

195

196
class FashionMNIST(MNIST):
197
198
199
    """`Fashion-MNIST <https://github.com/zalandoresearch/fashion-mnist>`_ Dataset.

    Args:
200
201
        root (string): Root directory of dataset where ``FashionMNIST/processed/training.pt``
            and  ``FashionMNIST/processed/test.pt`` exist.
202
203
204
205
206
207
208
209
210
        train (bool, optional): If True, creates dataset from ``training.pt``,
            otherwise from ``test.pt``.
        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.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
211
    """
212
213
214
215
    mirrors = [
        "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/"
    ]

216
    resources = [
217
218
219
220
        ("train-images-idx3-ubyte.gz", "8d4fb7e6c68d591d4c3dfef9ec88bf0d"),
        ("train-labels-idx1-ubyte.gz", "25c81989df183df01b3e8a0aad5dffbe"),
        ("t10k-images-idx3-ubyte.gz", "bef4ecab320f06d8554ea6380940ec79"),
        ("t10k-labels-idx1-ubyte.gz", "bb300cfdad3c16e7a12a480ee83cd310")
221
    ]
222
223
    classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal',
               'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
224
225


hysts's avatar
hysts committed
226
227
228
229
class KMNIST(MNIST):
    """`Kuzushiji-MNIST <https://github.com/rois-codh/kmnist>`_ Dataset.

    Args:
230
231
        root (string): Root directory of dataset where ``KMNIST/processed/training.pt``
            and  ``KMNIST/processed/test.pt`` exist.
hysts's avatar
hysts committed
232
233
234
235
236
237
238
239
240
241
        train (bool, optional): If True, creates dataset from ``training.pt``,
            otherwise from ``test.pt``.
        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.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
    """
242
243
244
245
    mirrors = [
        "http://codh.rois.ac.jp/kmnist/dataset/kmnist/"
    ]

246
    resources = [
247
248
249
250
        ("train-images-idx3-ubyte.gz", "bdb82020997e1d708af4cf47b453dcf7"),
        ("train-labels-idx1-ubyte.gz", "e144d726b3acfaa3e44228e80efcd344"),
        ("t10k-images-idx3-ubyte.gz", "5c965bf0a639b31b8f53240b1b52f4d7"),
        ("t10k-labels-idx1-ubyte.gz", "7320c461ea6c1c855c0b718fb2a4b134")
hysts's avatar
hysts committed
251
252
253
254
    ]
    classes = ['o', 'ki', 'su', 'tsu', 'na', 'ha', 'ma', 'ya', 're', 'wo']


255
class EMNIST(MNIST):
Alex Alemi's avatar
Alex Alemi committed
256
    """`EMNIST <https://www.westernsydney.edu.au/bens/home/reproducible_research/emnist>`_ Dataset.
257
258

    Args:
259
260
        root (string): Root directory of dataset where ``EMNIST/processed/training.pt``
            and  ``EMNIST/processed/test.pt`` exist.
261
262
263
264
265
266
267
268
269
270
271
272
273
        split (string): The dataset has 6 different splits: ``byclass``, ``bymerge``,
            ``balanced``, ``letters``, ``digits`` and ``mnist``. This argument specifies
            which one to use.
        train (bool, optional): If True, creates dataset from ``training.pt``,
            otherwise from ``test.pt``.
        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.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
    """
Philip Meier's avatar
Philip Meier committed
274
    url = 'https://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/gzip.zip'
275
    md5 = "58c8d27c78d21e728a6bc7b3cc06412e"
276
    splits = ('byclass', 'bymerge', 'balanced', 'letters', 'digits', 'mnist')
277
    # Merged Classes assumes Same structure for both uppercase and lowercase version
278
279
    _merged_classes = {'c', 'i', 'j', 'k', 'l', 'm', 'o', 'p', 's', 'u', 'v', 'w', 'x', 'y', 'z'}
    _all_classes = set(string.digits + string.ascii_letters)
280
    classes_split_dict = {
281
        'byclass': sorted(list(_all_classes)),
282
283
        'bymerge': sorted(list(_all_classes - _merged_classes)),
        'balanced': sorted(list(_all_classes - _merged_classes)),
284
        'letters': ['N/A'] + list(string.ascii_lowercase),
285
286
287
        'digits': list(string.digits),
        'mnist': list(string.digits),
    }
288

289
    def __init__(self, root: str, split: str, **kwargs: Any) -> None:
290
        self.split = verify_str_arg(split, "split", self.splits)
291
292
293
        self.training_file = self._training_file(split)
        self.test_file = self._test_file(split)
        super(EMNIST, self).__init__(root, **kwargs)
294
        self.classes = self.classes_split_dict[self.split]
Tian Qi Chen's avatar
Tian Qi Chen committed
295

296
    @staticmethod
297
    def _training_file(split) -> str:
298
299
        return 'training_{}.pt'.format(split)

300
    @staticmethod
301
    def _test_file(split) -> str:
302
303
        return 'test_{}.pt'.format(split)

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
    @property
    def _file_prefix(self) -> str:
        return f"emnist-{self.split}-{'train' if self.train else 'test'}"

    @property
    def images_file(self) -> str:
        return os.path.join(self.raw_folder, f"{self._file_prefix}-images-idx3-ubyte")

    @property
    def labels_file(self) -> str:
        return os.path.join(self.raw_folder, f"{self._file_prefix}-labels-idx1-ubyte")

    def _load_data(self):
        return read_image_file(self.images_file), read_label_file(self.labels_file)

    def _check_exists(self) -> bool:
        return all(check_integrity(file) for file in (self.images_file, self.labels_file))

322
    def download(self) -> None:
323
        """Download the EMNIST data if it doesn't exist already."""
324

325
326
327
        if self._check_exists():
            return

328
        os.makedirs(self.raw_folder, exist_ok=True)
329

330
        download_and_extract_archive(self.url, download_root=self.raw_folder, md5=self.md5)
331
        gzip_folder = os.path.join(self.raw_folder, 'gzip')
332
333
        for gzip_file in os.listdir(gzip_folder):
            if gzip_file.endswith('.gz'):
334
                extract_archive(os.path.join(gzip_folder, gzip_file), self.raw_folder)
335
        shutil.rmtree(gzip_folder)
336
337


338
339
340
341
class QMNIST(MNIST):
    """`QMNIST <https://github.com/facebookresearch/qmnist>`_ Dataset.

    Args:
342
        root (string): Root directory of dataset whose ``processed``
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
            subdir contains torch binary files with the datasets.
        what (string,optional): Can be 'train', 'test', 'test10k',
            'test50k', or 'nist' for respectively the mnist compatible
            training set, the 60k qmnist testing set, the 10k qmnist
            examples that match the mnist testing set, the 50k
            remaining qmnist testing examples, or all the nist
            digits. The default is to select 'train' or 'test'
            according to the compatibility argument 'train'.
        compat (bool,optional): A boolean that says whether the target
            for each example is class number (for compatibility with
            the MNIST dataloader) or a torch vector containing the
            full qmnist information. Default=True.
        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.
        transform (callable, optional): A function/transform that
            takes in an PIL image and returns a transformed
            version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform
            that takes in the target and transforms it.
        train (bool,optional,compatibility): When argument 'what' is
            not specified, this boolean decides whether to load the
            training set ot the testing set.  Default: True.
    """

    subsets = {
        'train': 'train',
370
371
372
        'test': 'test',
        'test10k': 'test',
        'test50k': 'test',
373
374
        'nist': 'nist'
    }
375
    resources: Dict[str, List[Tuple[str, str]]] = {  # type: ignore[assignment]
376
377
378
379
380
381
382
383
384
385
386
387
        'train': [('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-images-idx3-ubyte.gz',
                   'ed72d4157d28c017586c42bc6afe6370'),
                  ('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-labels-idx2-int.gz',
                   '0058f8dd561b90ffdd0f734c6a30e5e4')],
        'test': [('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-images-idx3-ubyte.gz',
                  '1394631089c404de565df7b7aeaf9412'),
                 ('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-labels-idx2-int.gz',
                  '5b5b05890a5e13444e108efe57b788aa')],
        'nist': [('https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-images-idx3-ubyte.xz',
                  '7f124b3b8ab81486c9d8c2749c17f834'),
                 ('https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-labels-idx2-int.xz',
                  '5ed0e788978e45d4a8bd4b7caec3d79d')]
388
389
390
391
    }
    classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
               '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']

392
393
394
395
    def __init__(
            self, root: str, what: Optional[str] = None, compat: bool = True,
            train: bool = True, **kwargs: Any
    ) -> None:
396
397
        if what is None:
            what = 'train' if train else 'test'
398
        self.what = verify_str_arg(what, "what", tuple(self.subsets.keys()))
399
400
401
402
403
404
        self.compat = compat
        self.data_file = what + '.pt'
        self.training_file = self.data_file
        self.test_file = self.data_file
        super(QMNIST, self).__init__(root, train, **kwargs)

405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
    @property
    def images_file(self) -> str:
        (url, _), _ = self.resources[self.subsets[self.what]]
        return os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0])

    @property
    def labels_file(self) -> str:
        _, (url, _) = self.resources[self.subsets[self.what]]
        return os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0])

    def _check_exists(self) -> bool:
        return all(check_integrity(file) for file in (self.images_file, self.labels_file))

    def _load_data(self):
        data = read_sn3_pascalvincent_tensor(self.images_file)
        assert (data.dtype == torch.uint8)
        assert (data.ndimension() == 3)

        targets = read_sn3_pascalvincent_tensor(self.labels_file).long()
        assert (targets.ndimension() == 2)

        if self.what == 'test10k':
            data = data[0:10000, :, :].clone()
            targets = targets[0:10000, :].clone()
        elif self.what == 'test50k':
            data = data[10000:, :, :].clone()
            targets = targets[10000:, :].clone()

        return data, targets

435
    def download(self) -> None:
436
        """Download the QMNIST data if it doesn't exist already.
437
438
439
440
           Note that we only download what has been asked for (argument 'what').
        """
        if self._check_exists():
            return
441

442
        os.makedirs(self.raw_folder, exist_ok=True)
443
        split = self.resources[self.subsets[self.what]]
444

445
        for url, md5 in split:
446
447
448
            filename = url.rpartition('/')[2]
            file_path = os.path.join(self.raw_folder, filename)
            if not os.path.isfile(file_path):
449
                download_and_extract_archive(url, self.raw_folder, filename=filename, md5=md5)
450

451
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
452
453
454
455
456
457
458
459
460
461
462
        # redefined to handle the compat flag
        img, target = self.data[index], self.targets[index]
        img = Image.fromarray(img.numpy(), mode='L')
        if self.transform is not None:
            img = self.transform(img)
        if self.compat:
            target = int(target[0])
        if self.target_transform is not None:
            target = self.target_transform(target)
        return img, target

463
    def extra_repr(self) -> str:
464
465
466
        return "Split: {}".format(self.what)


467
def get_int(b: bytes) -> int:
468
    return int(codecs.encode(b, 'hex'), 16)
Tian Qi Chen's avatar
Tian Qi Chen committed
469

470

471
472
473
474
475
476
477
478
479
480
SN3_PASCALVINCENT_TYPEMAP = {
    8: (torch.uint8, np.uint8, np.uint8),
    9: (torch.int8, np.int8, np.int8),
    11: (torch.int16, np.dtype('>i2'), 'i2'),
    12: (torch.int32, np.dtype('>i4'), 'i4'),
    13: (torch.float32, np.dtype('>f4'), 'f4'),
    14: (torch.float64, np.dtype('>f8'), 'f8')
}


481
def read_sn3_pascalvincent_tensor(path: str, strict: bool = True) -> torch.Tensor:
482
483
484
485
    """Read a SN3 file in "Pascal Vincent" format (Lush file 'libidx/idx-io.lsh').
       Argument may be a filename, compressed filename, or file object.
    """
    # read
486
    with open(path, "rb") as f:
487
488
489
490
491
        data = f.read()
    # parse
    magic = get_int(data[0:4])
    nd = magic % 256
    ty = magic // 256
492
493
    assert 1 <= nd <= 3
    assert 8 <= ty <= 14
494
    m = SN3_PASCALVINCENT_TYPEMAP[ty]
495
496
497
    s = [get_int(data[4 * (i + 1): 4 * (i + 2)]) for i in range(nd)]
    parsed = np.frombuffer(data, dtype=m[1], offset=(4 * (nd + 1)))
    assert parsed.shape[0] == np.prod(s) or not strict
498
    return torch.from_numpy(parsed.astype(m[2])).view(*s)
499
500


501
def read_label_file(path: str) -> torch.Tensor:
502
    x = read_sn3_pascalvincent_tensor(path, strict=False)
503
504
505
    assert(x.dtype == torch.uint8)
    assert(x.ndimension() == 1)
    return x.long()
Tian Qi Chen's avatar
Tian Qi Chen committed
506

507

508
def read_image_file(path: str) -> torch.Tensor:
509
    x = read_sn3_pascalvincent_tensor(path, strict=False)
510
511
512
    assert(x.dtype == torch.uint8)
    assert(x.ndimension() == 3)
    return x