-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
77 lines (58 loc) · 2.11 KB
/
main.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
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
import gzip
import numpy as np
from numpy import ndarray
from tqdm import tqdm
from sklearn.metrics import f1_score
import numpytorch as npt
from numpytorch import tensor, nn, optim
from assignment import MNISTClassificationModel
def get_mnist() -> tuple[ndarray, ndarray, ndarray, ndarray]:
path_x = 'data/train-images-idx3-ubyte.gz'
path_y = 'data/train-labels-idx1-ubyte.gz'
with gzip.open(path_x, 'rb') as f:
images = np.frombuffer(f.read(), np.uint8, offset=16)
x = images.reshape(-1, 1, 28, 28).astype(np.float32) / 255.
with gzip.open(path_y, 'rb') as f:
labels = np.frombuffer(f.read(), np.uint8, offset=8)
y = labels
x_train, x_val = x[:50000], x[50000:]
y_train, y_val = y[:50000], y[50000:]
return (x_train, y_train, x_val, y_val)
def val(model: nn.Module, x_val: ndarray, y_val: ndarray) -> tuple[float, float]:
preds = []
for i in range(0, len(x_val), 32):
x = tensor(x_val[i:i+32])
preds += model(x).arr.argmax(-1).tolist()
macro = f1_score(y_val, preds, average="macro")
micro = f1_score(y_val, preds, average="micro")
return macro, micro
if __name__ == '__main__':
x_train, y_train, x_val, y_val = get_mnist()
n_batch = 32
n_iter = 50000
n_print = 100
n_val = 2000
lr = 1e-03
model = MNISTClassificationModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr)
buf = 0
for i in tqdm(range(1, n_iter+1)):
idx = np.random.permutation(50000)[:n_batch]
x, y = tensor(x_train[idx]), tensor(y_train[idx])
optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, y)
if np.isnan(loss.item()):
break
loss.backward()
optimizer.step()
buf += loss.item()
if i % n_print == 0:
print(buf / n_print)
buf = 0
if i % n_val == 0:
macro, micro = val(model, x_val, y_val)
print(f"macro: {macro:.9f} micro: {micro:.9f}")
macro, micro = val(model, x_val, y_val)
print(f"\nfinal score\nmacro: {macro:.9f} micro: {micro:.9f}")