-
Notifications
You must be signed in to change notification settings - Fork 0
/
mnist.py
52 lines (41 loc) · 1.47 KB
/
mnist.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import numpy as np
from urllib import request
import gzip
import pickle
import os
filename = [
["training_images", "train-images-idx3-ubyte.gz",
"mnist/train-images-idx3-ubyte.gz"],
["test_images", "t10k-images-idx3-ubyte.gz", "mnist/t10k-images-idx3-ubyte.gz"],
["training_labels", "train-labels-idx1-ubyte.gz",
"mnist/train-labels-idx1-ubyte.gz"],
["test_labels", "t10k-labels-idx1-ubyte.gz", "mnist/t10k-labels-idx1-ubyte.gz"]
]
def download_mnist():
base_url = "http://yann.lecun.com/exdb/mnist/"
for name in filename:
print("Downloading " + name[0] + "...")
os.makedirs("mnist", exist_ok=True)
request.urlretrieve(base_url + name[1], name[2])
print("Download complete.")
def save_mnist():
mnist = {}
for name in filename[:2]:
with gzip.open(name[2], 'rb') as f:
mnist[name[0]] = np.frombuffer(
f.read(), np.uint8, offset=16).reshape(-1, 28*28)
for name in filename[-2:]:
with gzip.open(name[2], 'rb') as f:
mnist[name[0]] = np.frombuffer(f.read(), np.uint8, offset=8)
with open("mnist/mnist.pkl", 'wb') as f:
pickle.dump(mnist, f)
print("Save complete.")
def init():
download_mnist()
save_mnist()
def load():
with open("mnist/mnist.pkl", 'rb') as f:
mnist = pickle.load(f)
return mnist["training_images"], mnist["training_labels"], mnist["test_images"], mnist["test_labels"]
if __name__ == '__main__':
init()