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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
|
"""Download MNIST and prepare .npy files expected by main.py."""
import argparse
import gzip
import shutil
import urllib.request
from pathlib import Path
import idx2numpy
import numpy as np
BASE_URL = "https://storage.googleapis.com/cvdf-datasets/mnist"
FILES = {
"train-images-idx3-ubyte.gz": f"{BASE_URL}/train-images-idx3-ubyte.gz",
"train-labels-idx1-ubyte.gz": f"{BASE_URL}/train-labels-idx1-ubyte.gz",
"t10k-images-idx3-ubyte.gz": f"{BASE_URL}/t10k-images-idx3-ubyte.gz",
"t10k-labels-idx1-ubyte.gz": f"{BASE_URL}/t10k-labels-idx1-ubyte.gz",
}
def download(url: str, dest: Path) -> None:
if dest.exists():
print(f"[skip] {dest} already exists")
return
dest.parent.mkdir(parents=True, exist_ok=True)
print(f"[download] {url} -> {dest}", end="", flush=True)
urllib.request.urlretrieve(url, dest)
print(" OK")
def gunzip(src: Path, dst: Path) -> None:
if dst.exists():
print(f"[skip] {dst} already exists")
return
print(f"[gunzip] {src} -> {dst}", end="", flush=True)
with gzip.open(src, "rb") as f_in, open(dst, "wb") as f_out:
shutil.copyfileobj(f_in, f_out)
print(" OK")
def convert_dataset(images_path: Path, labels_path: Path, out_path: Path) -> None:
print(f"[convert] {images_path.name} + {labels_path.name} -> {out_path}")
images = idx2numpy.convert_from_file(str(images_path)) / 255.0
labels = idx2numpy.convert_from_file(str(labels_path))
new_labels = np.zeros((len(labels), 10))
new_labels[np.arange(len(labels)), labels] = 1
dataset = list(zip(images, new_labels, strict=False))
out_path.parent.mkdir(parents=True, exist_ok=True)
with open(out_path, "wb") as f:
np.save(f, np.array(dataset, dtype=object))
print(f" saved {len(dataset)} samples to {out_path}")
def ensure_mnist_weights(path: Path, layer_sizes=(784, 32, 16, 10)) -> None:
if path.exists():
print(f"[skip] {path} already exists")
return
print(f"[init] creating {path} with layers {layer_sizes}")
from MultiLayerPerceptron.main import MultiLayerPerceptron
mlp = MultiLayerPerceptron(list(layer_sizes))
mlp.save(str(path))
print(" OK")
def main() -> None:
parser = argparse.ArgumentParser(description="Download MNIST and prepare npy files")
parser.add_argument("--raw-dir", default="data/raw", help="directory for raw idx files")
parser.add_argument("--force", action="store_true", help="re-download even if exists")
args = parser.parse_args()
raw_dir = Path(args.raw_dir)
if args.force:
for name in FILES:
p = raw_dir / name
if p.exists():
p.unlink()
ungz = raw_dir / name.replace(".gz", "")
if ungz.exists():
ungz.unlink()
# 1. download
for name, url in FILES.items():
download(url, raw_dir / name)
# 2. gunzip
for name in FILES:
gunzip(raw_dir / name, raw_dir / name.replace(".gz", ""))
# 3. convert to npy expected by main.py
convert_dataset(
raw_dir / "train-images-idx3-ubyte",
raw_dir / "train-labels-idx1-ubyte",
Path("data/train.npy"),
)
convert_dataset(
raw_dir / "t10k-images-idx3-ubyte",
raw_dir / "t10k-labels-idx1-ubyte",
Path("data/test.npy"),
)
# 4. ensure initial weights
ensure_mnist_weights(Path("mnist.npy"))
print("[done] all files ready")
if __name__ == "__main__":
main()
|
