mnist.py 20.8 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
11
import gzip
import lzma
12
from typing import Any, Callable, Dict, IO, List, Optional, Tuple, Union
13
from .utils import download_url, download_and_extract_archive, extract_archive, \
14
    verify_str_arg
Tian Qi Chen's avatar
Tian Qi Chen committed
15

16

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

    Args:
21
22
        root (string): Root directory of dataset where ``MNIST/processed/training.pt``
            and  ``MNIST/processed/test.pt`` exist.
23
24
25
26
27
28
29
30
31
32
        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.
    """
33
34
35
36
37
38

    resources = [
        ("http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
        ("http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
        ("http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
        ("http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
Tian Qi Chen's avatar
Tian Qi Chen committed
39
    ]
40

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

46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
    @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

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

        if download:
            self.download()

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

        if self.train:
86
            data_file = self.training_file
Tian Qi Chen's avatar
Tian Qi Chen committed
87
        else:
88
89
            data_file = self.test_file
        self.data, self.targets = torch.load(os.path.join(self.processed_folder, data_file))
Tian Qi Chen's avatar
Tian Qi Chen committed
90

91
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
92
93
94
95
96
97
98
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
99
        img, target = self.data[index], int(self.targets[index])
Tian Qi Chen's avatar
Tian Qi Chen committed
100
101
102
103
104
105
106
107
108
109
110
111
112

        # 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

113
    def __len__(self) -> int:
114
        return len(self.data)
Tian Qi Chen's avatar
Tian Qi Chen committed
115

116
    @property
117
    def raw_folder(self) -> str:
118
119
120
        return os.path.join(self.root, self.__class__.__name__, 'raw')

    @property
121
    def processed_folder(self) -> str:
122
123
124
        return os.path.join(self.root, self.__class__.__name__, 'processed')

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

128
    def _check_exists(self) -> bool:
129
130
131
132
        return (os.path.exists(os.path.join(self.processed_folder,
                                            self.training_file)) and
                os.path.exists(os.path.join(self.processed_folder,
                                            self.test_file)))
133

134
    def download(self) -> None:
135
        """Download the MNIST data if it doesn't exist in processed_folder already."""
Tian Qi Chen's avatar
Tian Qi Chen committed
136
137
138
139

        if self._check_exists():
            return

140
141
        os.makedirs(self.raw_folder, exist_ok=True)
        os.makedirs(self.processed_folder, exist_ok=True)
Tian Qi Chen's avatar
Tian Qi Chen committed
142

143
        # download files
144
        for url, md5 in self.resources:
Tian Qi Chen's avatar
Tian Qi Chen committed
145
            filename = url.rpartition('/')[2]
146
            download_and_extract_archive(url, download_root=self.raw_folder, filename=filename, md5=md5)
Tian Qi Chen's avatar
Tian Qi Chen committed
147
148

        # process and save as torch files
Adam Paszke's avatar
Adam Paszke committed
149
150
        print('Processing...')

Tian Qi Chen's avatar
Tian Qi Chen committed
151
        training_set = (
152
153
            read_image_file(os.path.join(self.raw_folder, 'train-images-idx3-ubyte')),
            read_label_file(os.path.join(self.raw_folder, 'train-labels-idx1-ubyte'))
Tian Qi Chen's avatar
Tian Qi Chen committed
154
155
        )
        test_set = (
156
157
            read_image_file(os.path.join(self.raw_folder, 't10k-images-idx3-ubyte')),
            read_label_file(os.path.join(self.raw_folder, 't10k-labels-idx1-ubyte'))
Tian Qi Chen's avatar
Tian Qi Chen committed
158
        )
159
        with open(os.path.join(self.processed_folder, self.training_file), 'wb') as f:
Tian Qi Chen's avatar
Tian Qi Chen committed
160
            torch.save(training_set, f)
161
        with open(os.path.join(self.processed_folder, self.test_file), 'wb') as f:
Tian Qi Chen's avatar
Tian Qi Chen committed
162
163
164
165
            torch.save(test_set, f)

        print('Done!')

166
    def extra_repr(self) -> str:
167
        return "Split: {}".format("Train" if self.train is True else "Test")
168

169

170
class FashionMNIST(MNIST):
171
172
173
    """`Fashion-MNIST <https://github.com/zalandoresearch/fashion-mnist>`_ Dataset.

    Args:
174
175
        root (string): Root directory of dataset where ``FashionMNIST/processed/training.pt``
            and  ``FashionMNIST/processed/test.pt`` exist.
176
177
178
179
180
181
182
183
184
        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.
185
    """
186
187
188
189
190
191
192
193
194
    resources = [
        ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz",
         "8d4fb7e6c68d591d4c3dfef9ec88bf0d"),
        ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz",
         "25c81989df183df01b3e8a0aad5dffbe"),
        ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz",
         "bef4ecab320f06d8554ea6380940ec79"),
        ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz",
         "bb300cfdad3c16e7a12a480ee83cd310")
195
    ]
196
197
    classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal',
               'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
198
199


hysts's avatar
hysts committed
200
201
202
203
class KMNIST(MNIST):
    """`Kuzushiji-MNIST <https://github.com/rois-codh/kmnist>`_ Dataset.

    Args:
204
205
        root (string): Root directory of dataset where ``KMNIST/processed/training.pt``
            and  ``KMNIST/processed/test.pt`` exist.
hysts's avatar
hysts committed
206
207
208
209
210
211
212
213
214
215
        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.
    """
216
217
218
219
220
    resources = [
        ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/train-images-idx3-ubyte.gz", "bdb82020997e1d708af4cf47b453dcf7"),
        ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/train-labels-idx1-ubyte.gz", "e144d726b3acfaa3e44228e80efcd344"),
        ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/t10k-images-idx3-ubyte.gz", "5c965bf0a639b31b8f53240b1b52f4d7"),
        ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/t10k-labels-idx1-ubyte.gz", "7320c461ea6c1c855c0b718fb2a4b134")
hysts's avatar
hysts committed
221
222
223
224
    ]
    classes = ['o', 'ki', 'su', 'tsu', 'na', 'ha', 'ma', 'ya', 're', 'wo']


225
class EMNIST(MNIST):
Alex Alemi's avatar
Alex Alemi committed
226
    """`EMNIST <https://www.westernsydney.edu.au/bens/home/reproducible_research/emnist>`_ Dataset.
227
228

    Args:
229
230
        root (string): Root directory of dataset where ``EMNIST/processed/training.pt``
            and  ``EMNIST/processed/test.pt`` exist.
231
232
233
234
235
236
237
238
239
240
241
242
243
        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
244
245
246
247
248
    # Updated URL from https://www.nist.gov/node/1298471/emnist-dataset since the
    # _official_ download link
    # https://cloudstor.aarnet.edu.au/plus/s/ZNmuFiuQTqZlu9W/download
    # is (currently) unavailable
    url = 'http://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/gzip.zip'
249
    md5 = "58c8d27c78d21e728a6bc7b3cc06412e"
250
    splits = ('byclass', 'bymerge', 'balanced', 'letters', 'digits', 'mnist')
251
    # Merged Classes assumes Same structure for both uppercase and lowercase version
252
253
    _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)
254
    classes_split_dict = {
255
        'byclass': sorted(list(_all_classes)),
256
257
        'bymerge': sorted(list(_all_classes - _merged_classes)),
        'balanced': sorted(list(_all_classes - _merged_classes)),
258
        'letters': ['N/A'] + list(string.ascii_lowercase),
259
260
261
        'digits': list(string.digits),
        'mnist': list(string.digits),
    }
262

263
    def __init__(self, root: str, split: str, **kwargs: Any) -> None:
264
        self.split = verify_str_arg(split, "split", self.splits)
265
266
267
        self.training_file = self._training_file(split)
        self.test_file = self._test_file(split)
        super(EMNIST, self).__init__(root, **kwargs)
268
        self.classes = self.classes_split_dict[self.split]
Tian Qi Chen's avatar
Tian Qi Chen committed
269

270
    @staticmethod
271
    def _training_file(split) -> str:
272
273
        return 'training_{}.pt'.format(split)

274
    @staticmethod
275
    def _test_file(split) -> str:
276
277
        return 'test_{}.pt'.format(split)

278
    def download(self) -> None:
279
280
        """Download the EMNIST data if it doesn't exist in processed_folder already."""
        import shutil
281

282
283
284
        if self._check_exists():
            return

285
286
        os.makedirs(self.raw_folder, exist_ok=True)
        os.makedirs(self.processed_folder, exist_ok=True)
287

288
        # download files
289
        print('Downloading and extracting zip archive')
290
        download_and_extract_archive(self.url, download_root=self.raw_folder, filename="emnist.zip",
291
                                     remove_finished=True, md5=self.md5)
292
        gzip_folder = os.path.join(self.raw_folder, 'gzip')
293
294
        for gzip_file in os.listdir(gzip_folder):
            if gzip_file.endswith('.gz'):
295
                extract_archive(os.path.join(gzip_folder, gzip_file), gzip_folder)
296
297
298
299
300

        # process and save as torch files
        for split in self.splits:
            print('Processing ' + split)
            training_set = (
301
302
                read_image_file(os.path.join(gzip_folder, 'emnist-{}-train-images-idx3-ubyte'.format(split))),
                read_label_file(os.path.join(gzip_folder, 'emnist-{}-train-labels-idx1-ubyte'.format(split)))
303
304
            )
            test_set = (
305
306
                read_image_file(os.path.join(gzip_folder, 'emnist-{}-test-images-idx3-ubyte'.format(split))),
                read_label_file(os.path.join(gzip_folder, 'emnist-{}-test-labels-idx1-ubyte'.format(split)))
307
            )
308
            with open(os.path.join(self.processed_folder, self._training_file(split)), 'wb') as f:
309
                torch.save(training_set, f)
310
            with open(os.path.join(self.processed_folder, self._test_file(split)), 'wb') as f:
311
                torch.save(test_set, f)
312
        shutil.rmtree(gzip_folder)
313
314
315
316

        print('Done!')


317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
class QMNIST(MNIST):
    """`QMNIST <https://github.com/facebookresearch/qmnist>`_ Dataset.

    Args:
        root (string): Root directory of dataset whose ``processed''
            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',
350
351
352
        'test': 'test',
        'test10k': 'test',
        'test50k': 'test',
353
354
        'nist': 'nist'
    }
355
    resources: Dict[str, List[Tuple[str, str]]] = {  # type: ignore[assignment]
356
357
358
359
360
361
362
363
364
365
366
367
        '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')]
368
369
370
371
    }
    classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
               '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']

372
373
374
375
    def __init__(
            self, root: str, what: Optional[str] = None, compat: bool = True,
            train: bool = True, **kwargs: Any
    ) -> None:
376
377
        if what is None:
            what = 'train' if train else 'test'
378
        self.what = verify_str_arg(what, "what", tuple(self.subsets.keys()))
379
380
381
382
383
384
        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)

385
    def download(self) -> None:
386
387
388
389
390
        """Download the QMNIST data if it doesn't exist in processed_folder already.
           Note that we only download what has been asked for (argument 'what').
        """
        if self._check_exists():
            return
391
392
        os.makedirs(self.raw_folder, exist_ok=True)
        os.makedirs(self.processed_folder, exist_ok=True)
393
        split = self.resources[self.subsets[self.what]]
394
395
396
        files = []

        # download data files if not already there
397
        for url, md5 in split:
398
399
400
            filename = url.rpartition('/')[2]
            file_path = os.path.join(self.raw_folder, filename)
            if not os.path.isfile(file_path):
401
                download_url(url, root=self.raw_folder, filename=filename, md5=md5)
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
            files.append(file_path)

        # process and save as torch files
        print('Processing...')
        data = read_sn3_pascalvincent_tensor(files[0])
        assert(data.dtype == torch.uint8)
        assert(data.ndimension() == 3)
        targets = read_sn3_pascalvincent_tensor(files[1]).long()
        assert(targets.ndimension() == 2)
        if self.what == 'test10k':
            data = data[0:10000, :, :].clone()
            targets = targets[0:10000, :].clone()
        if self.what == 'test50k':
            data = data[10000:, :, :].clone()
            targets = targets[10000:, :].clone()
        with open(os.path.join(self.processed_folder, self.data_file), 'wb') as f:
            torch.save((data, targets), f)

420
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
421
422
423
424
425
426
427
428
429
430
431
        # 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

432
    def extra_repr(self) -> str:
433
434
435
        return "Split: {}".format(self.what)


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

439

440
def open_maybe_compressed_file(path: Union[str, IO]) -> Union[IO, gzip.GzipFile]:
441
442
443
444
445
446
447
448
449
450
451
452
    """Return a file object that possibly decompresses 'path' on the fly.
       Decompression occurs when argument `path` is a string and ends with '.gz' or '.xz'.
    """
    if not isinstance(path, torch._six.string_classes):
        return path
    if path.endswith('.gz'):
        return gzip.open(path, 'rb')
    if path.endswith('.xz'):
        return lzma.open(path, 'rb')
    return open(path, 'rb')


453
454
455
456
457
458
459
460
461
462
463
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')
}


def read_sn3_pascalvincent_tensor(path: Union[str, IO], strict: bool = True) -> torch.Tensor:
464
465
466
467
468
469
470
471
472
473
474
475
    """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
    with open_maybe_compressed_file(path) as f:
        data = f.read()
    # parse
    magic = get_int(data[0:4])
    nd = magic % 256
    ty = magic // 256
    assert nd >= 1 and nd <= 3
    assert ty >= 8 and ty <= 14
476
    m = SN3_PASCALVINCENT_TYPEMAP[ty]
477
478
479
480
481
482
    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
    return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)


483
def read_label_file(path: str) -> torch.Tensor:
Tian Qi Chen's avatar
Tian Qi Chen committed
484
    with open(path, 'rb') as f:
485
486
487
488
        x = read_sn3_pascalvincent_tensor(f, strict=False)
    assert(x.dtype == torch.uint8)
    assert(x.ndimension() == 1)
    return x.long()
Tian Qi Chen's avatar
Tian Qi Chen committed
489

490

491
def read_image_file(path: str) -> torch.Tensor:
Tian Qi Chen's avatar
Tian Qi Chen committed
492
    with open(path, 'rb') as f:
493
494
495
496
        x = read_sn3_pascalvincent_tensor(f, strict=False)
    assert(x.dtype == torch.uint8)
    assert(x.ndimension() == 3)
    return x