1
Input (28×28×1)
2
    ↓
3
Conv2D (8 filters, 3×3 kernel)  → (26×26×8)
4
    ↓
5
ReLU
6
    ↓
7
Conv2D (16 filters, 3×3 kernel) → (24×24×16)
8
    ↓
9
ReLU
10
    ↓
11
MaxPool (2×2)                    → (12×12×16)
12
    ↓
13
Flatten                          → (2304,)
14
    ↓
15
Dense (10 neurons)               → (10,)
16
    ↓
17
Softmax

1
def conv2d_naive(X, W, b):
2
    B, H, W_in, C_in = X.shape
3
    out_c, _, k, _ = W.shape
4
    out_h = H - k + 1
5
    out_w = W_in - k + 1
6
    out = np.zeros((B, out_h, out_w, out_c))
7

8
    for b in range(B):           # batch
9
        for oc in range(out_c):  # output channel
10
            for i in range(out_h):  # height
11
                for j in range(out_w):  # width
12
                    for ic in range(C_in):  # input channel
13
                        for ki in range(k):  # kernel height
14
                            for kj in range(k):  # kernel width
15
                                out[b,i,j,oc] += X[b,i+ki,j+kj,ic] * W[oc,ic,ki,kj]
16
                    out[b,i,j,oc] += b[oc,0]
17
    return out

1
from numba import njit
2

3
@njit(cache=True, fastmath=True)
4
def my_function(x):
5
    # 你的代码
6
    return result

1
class MinimalCNNClassifier:
2
    def __init__(self, lr=0.01, use_adam=False):
3
        self.lr = lr
4
        self.use_adam = use_adam
5
        self.adam = TinyAdam(lr=lr) if use_adam else None
6

7
        # 网络层
8
        self.conv1 = Conv2D(1, 8, 3)      # 1→8 通道
9
        self.conv2 = Conv2D(8, 16, 3)     # 8→16 通道
10
        self.pool = MaxPool2x2()          # 2×2 池化
11
        self.fc = Dense(12 * 12 * 16, 10) # 全连接层
12

13
    def forward(self, X):
14
        # Conv1 + ReLU
15
        out = self.conv1.forward(X)
16
        out = relu(out)
17
        self.after_relu1 = out
18

19
        # Conv2 + ReLU
20
        out = self.conv2.forward(out)
21
        out = relu(out)
22
        self.after_relu2 = out
23

24
        # MaxPool
25
        out = self.pool.forward(out)
26

27
        # Flatten + Dense + Softmax
28
        out = flatten(out)
29
        out = self.fc.forward(out)
30
        self.y_pred = softmax(out)
31
        return self.y_pred
32

33
    def backward(self, y_true):
34
        # 输出层梯度
35
        grad = softmax_cross_entropy_grad(self.y_pred, y_true)
36

37
        # Dense 反向传播
38
        grad = self.fc.backward(grad)
39

40
        # Unflatten
41
        grad = grad.reshape(-1, 12, 12, 16)
42

43
        # MaxPool 反向传播
44
        grad = self.pool.backward(grad)
45

46
        # Conv2 反向传播
47
        grad = grad * relu_grad(self.after_relu2)
48
        grad = self.conv2.backward(grad)
49

50
        # Conv1 反向传播
51
        grad = grad * relu_grad(self.after_relu1)
52
        self.conv1.backward(grad)
53

54
    def step(self):
55
        if self.adam:
56
            self.adam.step([
57
                (self.fc.W, self.fc.dW),
58
                (self.fc.b, self.fc.db),
59
                (self.conv2.W, self.conv2.dW),
60
                (self.conv2.b, self.conv2.db),
61
                (self.conv1.W, self.conv1.dW),
62
                (self.conv1.b, self.conv1.db),
63
            ])
64
        else:
65
            self.fc.step(self.lr)
66
            self.conv2.step(self.lr)
67
            self.conv1.step(self.lr)

1
import numpy as np
2
import tqdm
3

4
np.random.seed(0)
5

6
# 加载数据
7
data_dir = "./mnist"
8
X_train, y_train, X_test, y_test = load_mnist_from_local(data_dir)
9

10
# 创建模型
11
model = MinimalCNNClassifier(lr=0.001, use_adam=True)
12

13
# 训练
14
def train(model, X_train, y_train, epochs, batch_size):
15
    loss_list = []
16
    for epoch in range(epochs):
17
        pbar = tqdm.tqdm(
18
            range(0, len(X_train), batch_size),
19
            desc=f"Epoch {epoch+1}/{epochs}"
20
        )
21
        for i in pbar:
22
            X_batch = X_train[i:i+batch_size]
23
            y_batch = y_train[i:i+batch_size]
24

25
            y_pred = model.forward(X_batch)
26
            loss = cross_entropy(y_pred, y_batch)
27
            loss_list.append(loss)
28

29
            model.backward(y_batch)
30
            model.step()
31

32
            pbar.set_postfix({"loss": f"{loss:.4f}"})
33
    return loss_list
34

35
loss_list = train(model, X_train, y_train, epochs=3, batch_size=64)
36

37
# 测试
38
def test(model, X_test, y_test):
39
    y_pred_test = model.forward(X_test)
40
    pred = np.argmax(y_pred_test, axis=1)
41
    true = np.argmax(y_test, axis=1)
42
    acc = (pred == true).mean()
43
    print(f"\nTest accuracy: {acc:.4f}")
44
    return acc
45

46
acc = test(model, X_test, y_test)

1
Epoch 1/3: 100%| 938/938 [00:54<00:00, 17.24it/s, loss=0.0286]
2
Epoch 2/3: 100%| 938/938 [00:54<00:00, 17.19it/s, loss=0.0267]
3
Epoch 3/3: 100%| 938/938 [00:49<00:00, 18.90it/s, loss=0.0153]
4

5
Test accuracy: 0.9806

1
class MinimalCNNClassifier:
2
    # ... 前面的代码 ...
3

4
    def save(self, path):
5
        np.savez(
6
            path,
7
            conv1_W=self.conv1.W,
8
            conv1_b=self.conv1.b,
9
            conv2_W=self.conv2.W,
10
            conv2_b=self.conv2.b,
11
            fc_W=self.fc.W,
12
            fc_b=self.fc.b,
13
        )
14
        print(f"Model saved to {path}")
15

16
    def load(self, path):
17
        data = np.load(path)
18
        self.conv1.W = data["conv1_W"]
19
        self.conv1.b = data["conv1_b"]
20
        self.conv2.W = data["conv2_W"]
21
        self.conv2.b = data["conv2_b"]
22
        self.fc.W = data["fc_W"]
23
        self.fc.b = data["fc_b"]
24
        print(f"Model loaded from {path}")

1
# 保存
2
model.save("models/cnn_model.npz")
3

4
# 加载
5
model2 = MinimalCNNClassifier(lr=0.001, use_adam=True)
6
model2.load("models/cnn_model.npz")
7

8
# 测试
9
acc = test(model2, X_test, y_test)  # 0.9806

1
Model saved to models/cnn_model.npz!
2
Reloading model...
3
Model loaded from models/cnn_model.npz
4
Reload complete.
5
Reloaded model accuracy: 0.9806

1
class MinimalCNNClassifier:
2
    # ... 前面的代码 ...
3

4
    def predict_batch(self, X):
5
        """批量预测"""
6
        y_pred = self.forward(X)
7
        return np.argmax(y_pred, axis=1)
8

9
    def __call__(self, x):
10
        """
11
        单样本推理
12
        接受 (28,28)、(28,28,1) 或 (1,28,28,1) 格式的输入
13
        """
14
        if x.ndim == 2:  # (H, W)
15
            x = x[None, :, :, None]
16
        elif x.ndim == 3:  # (H, W, C)
17
            x = x[None, :, :, :]
18
        return self.predict_batch(x)[0]

1
# 单样本推理
2
sample = X_test[0].squeeze()  # (28, 28)
3
pred = model(sample)
4
true = np.argmax(y_test[0])
5
print(f"Predicted: {pred}, True: {true}")

1
Single sample predicted class: 7, true: 7

1
import matplotlib.pyplot as plt
2

3
# 第一层卷积核 (8, 1, 3, 3)
4
kernels = model.conv1.W.squeeze()  # (8, 3, 3)
5

6
plt.figure(figsize=(12, 2))
7
for i in range(8):
8
    plt.subplot(1, 8, i+1)
9
    plt.imshow(kernels[i], cmap='gray')
10
    plt.title(f"Filter {i+1}")
11
    plt.axis('off')
12
plt.tight_layout()
13
plt.show()

1
import numpy as np
2
import tqdm
3
import matplotlib.pyplot as plt
4
import os
5
import gzip
6
from numba import njit
7

8
# ============================================================
9
#         Load MNIST from Kaggle idx files (local)
10
# ============================================================
11
def _open_maybe_gz(path):
12
    if os.path.exists(path):
13
        return open(path, "rb")
14
    if os.path.exists(path + ".gz"):
15
        return gzip.open(path + ".gz", "rb")
16
    raise FileNotFoundError(f"Cannot find {path} or {path+'.gz'}")
17

18
def load_mnist_from_local(data_dir):
19
    train_images_path = os.path.join(data_dir, "train-images-idx3-ubyte")
20
    train_labels_path = os.path.join(data_dir, "train-labels-idx1-ubyte")
21
    test_images_path  = os.path.join(data_dir, "t10k-images-idx3-ubyte")
22
    test_labels_path  = os.path.join(data_dir, "t10k-labels-idx1-ubyte")
23

24
    # images: 16-byte header, then uint8 pixels
25
    with _open_maybe_gz(train_images_path) as f:
26
        data = np.frombuffer(f.read(), dtype=np.uint8, offset=16)
27
    X_train = data.reshape(-1, 28, 28, 1) / 255.0
28

29
    with _open_maybe_gz(test_images_path) as f:
30
        data = np.frombuffer(f.read(), dtype=np.uint8, offset=16)
31
    X_test = data.reshape(-1, 28, 28, 1) / 255.0
32

33
    # labels: 8-byte header, then uint8 labels
34
    with _open_maybe_gz(train_labels_path) as f:
35
        labels_train = np.frombuffer(f.read(), dtype=np.uint8, offset=8)
36
    with _open_maybe_gz(test_labels_path) as f:
37
        labels_test = np.frombuffer(f.read(), dtype=np.uint8, offset=8)
38

39
    y_train = np.zeros((labels_train.size, 10))
40
    y_train[np.arange(labels_train.size), labels_train] = 1
41

42
    y_test = np.zeros((labels_test.size, 10))
43
    y_test[np.arange(labels_test.size), labels_test] = 1
44

45
    return X_train, y_train, X_test, y_test
46

47
# ============================================================
48
#               Basic ops
49
# ============================================================
50
def relu(x):
51
    return np.maximum(0, x)
52

53
def relu_grad(x):
54
    return (x > 0).astype(float)
55

56
def softmax(x):
57
    x_max = np.max(x, axis=1, keepdims=True)
58
    ex = np.exp(x - x_max)
59
    return ex / np.sum(ex, axis=1, keepdims=True)
60

61
def cross_entropy(y_pred, y_true):
62
    eps = 1e-15
63
    y_pred = np.clip(y_pred, eps, 1 - eps)
64
    ce = -np.sum(y_true * np.log(y_pred), axis=1)
65
    return np.mean(ce)
66

67
def softmax_cross_entropy_grad(y_pred, y_true):
68
    return (y_pred - y_true) / y_true.shape[0]
69

70
def flatten(x):
71
    return x.reshape(x.shape[0], -1)
72

73
# ============================================================
74
#                Numba-accelerated convolution
75
# ============================================================
76
@njit(cache=True, fastmath=True)
77
def conv_forward_jit(X, W, b, k):
78
    B, H, W_in, C = X.shape
79
    out_c = W.shape[0]
80
    out_h = H - k + 1
81
    out_w = W_in - k + 1
82
    out = np.zeros((B, out_h, out_w, out_c), dtype=X.dtype)
83

84
    for b_idx in range(B):
85
        for oc in range(out_c):
86
            for i in range(out_h):
87
                for j in range(out_w):
88
                    acc = 0.0
89
                    for ic in range(C):
90
                        for ki in range(k):
91
                            for kj in range(k):
92
                                acc += X[b_idx, i + ki, j + kj, ic] * W[oc, ic, ki, kj]
93
                    out[b_idx, i, j, oc] = acc + b[oc, 0]
94
    return out
95

96
@njit(cache=True, fastmath=True)
97
def conv_backward_jit(X, grad, W, k):
98
    B, H, W_out, out_c = grad.shape
99
    _, XH, XW, in_c = X.shape
100

101
    dW = np.zeros_like(W)
102
    db = np.zeros((out_c, 1), dtype=grad.dtype)
103
    dX = np.zeros_like(X)
104

105
    for b_idx in range(B):
106
        for oc in range(out_c):
107
            for ic in range(in_c):
108
                for i in range(H):
109
                    for j in range(W_out):
110
                        g = grad[b_idx, i, j, oc]
111
                        if g == 0.0:
112
                            continue
113
                        for ki in range(k):
114
                            for kj in range(k):
115
                                dW[oc, ic, ki, kj] += g * X[b_idx, i + ki, j + kj, ic]
116
                                dX[b_idx, i + ki, j + kj, ic] += g * W[oc, ic, ki, kj]
117
            db[oc, 0] += np.sum(grad[b_idx, :, :, oc])
118
    return dW, db, dX
119

120
# ============================================================
121
#                Numba-accelerated max pooling
122
# ============================================================
123
@njit(cache=True, fastmath=True)
124
def maxpool_forward_jit(X):
125
    B, H, W, C = X.shape
126
    out = np.zeros((B, H // 2, W // 2, C), dtype=X.dtype)
127
    argmax = np.zeros((B, H // 2, W // 2, C), dtype=np.int64)
128

129
    for b_idx in range(B):
130
        for c in range(C):
131
            for i in range(0, H, 2):
132
                for j in range(0, W, 2):
133
                    max_val = X[b_idx, i, j, c]
134
                    max_idx = 0
135
                    idx = 1
136
                    v = X[b_idx, i, j + 1, c]
137
                    if v > max_val:
138
                        max_val = v
139
                        max_idx = idx
140
                    idx += 1
141
                    v = X[b_idx, i + 1, j, c]
142
                    if v > max_val:
143
                        max_val = v
144
                        max_idx = idx
145
                    idx += 1
146
                    v = X[b_idx, i + 1, j + 1, c]
147
                    if v > max_val:
148
                        max_val = v
149
                        max_idx = idx
150
                    out[b_idx, i // 2, j // 2, c] = max_val
151
                    argmax[b_idx, i // 2, j // 2, c] = max_idx
152
    return out, argmax
153

154
@njit(cache=True, fastmath=True)
155
def maxpool_backward_jit(grad, argmax, H_in, W_in):
156
    B, H2, W2, C = grad.shape
157
    dX = np.zeros((B, H_in, W_in, C), dtype=grad.dtype)
158

159
    for b_idx in range(B):
160
        for c in range(C):
161
            for i in range(H2):
162
                for j in range(W2):
163
                    idx = argmax[b_idx, i, j, c]
164
                    bi = idx // 2
165
                    bj = idx - bi * 2
166
                    dX[b_idx, i * 2 + bi, j * 2 + bj, c] = grad[b_idx, i, j, c]
167
    return dX
168

169
# ============================================================
170
#                     Convolution Layer
171
# ============================================================
172
class Conv2D:
173
    def __init__(self, in_channels, out_channels, kernel_size):
174
        self.in_c = in_channels
175
        self.out_c = out_channels
176
        self.k = kernel_size
177
        self.W = np.random.randn(out_channels, in_channels, kernel_size, kernel_size) * 0.1
178
        self.b = np.zeros((out_channels, 1))
179

180
    def forward(self, X):
181
        self.X = X
182
        self.out = conv_forward_jit(X, self.W, self.b, self.k)
183
        return self.out
184

185
    def backward(self, grad):
186
        self.dW, self.db, self.dX = conv_backward_jit(self.X, grad, self.W, self.k)
187
        return self.dX
188

189
    def step(self, lr):
190
        self.W -= lr * self.dW
191
        self.b -= lr * self.db
192

193
# ============================================================
194
#                     Max Pooling Layer
195
# ============================================================
196
class MaxPool2x2:
197
    def forward(self, X):
198
        self.X = X
199
        out, self.argmax = maxpool_forward_jit(X)
200
        return out
201

202
    def backward(self, grad):
203
        B, H, W, C = self.X.shape
204
        return maxpool_backward_jit(grad, self.argmax, H, W)
205

206
# ============================================================
207
#                     Dense Layer
208
# ============================================================
209
class Dense:
210
    def __init__(self, in_dim, out_dim):
211
        self.W = np.random.randn(in_dim, out_dim) * 0.01
212
        self.b = np.zeros((1, out_dim))
213

214
    def forward(self, X):
215
        self.X = X
216
        self.out = X @ self.W + self.b
217
        return self.out
218

219
    def backward(self, grad):
220
        self.dW = self.X.T @ grad
221
        self.db = np.sum(grad, axis=0, keepdims=True)
222
        return grad @ self.W.T
223

224
    def step(self, lr):
225
        self.W -= lr * self.dW
226
        self.b -= lr * self.db
227

228
# ============================================================
229
#                 CNN Classifier
230
# ============================================================
231
class MinimalCNNClassifier:
232
    def __init__(self, lr=0.01, use_adam=False, hidden_dim=44):
233
        self.lr = lr
234
        self.use_adam = use_adam
235
        self.adam = TinyAdam(lr=lr) if use_adam else None
236
        self.conv1 = Conv2D(1, 8, 3)
237
        self.conv2 = Conv2D(8, 16, 3)
238
        self.pool = MaxPool2x2()
239
        self.hidden_dim = hidden_dim  # choose 44 -> ~103K params total
240
        self.fc1 = Dense(12 * 12 * 16, hidden_dim)
241
        self.fc2 = Dense(hidden_dim, 10)
242

243
    def forward(self, X):
244
        out = self.conv1.forward(X)
245
        out = relu(out)
246
        self.after_relu1 = out
247

248
        out = self.conv2.forward(out)
249
        out = relu(out)
250
        self.after_relu2 = out
251

252
        out = self.pool.forward(out)
253
        out = flatten(out)
254

255
        out = self.fc1.forward(out)
256
        self.fc1_pre_relu = out
257
        out = relu(out)
258
        self.fc1_post_relu = out
259

260
        out = self.fc2.forward(out)
261
        self.y_pred = softmax(out)
262
        return self.y_pred
263

264
    def backward(self, y_true):
265
        grad = softmax_cross_entropy_grad(self.y_pred, y_true)
266

267
        grad = self.fc2.backward(grad)
268

269
        grad = grad * relu_grad(self.fc1_pre_relu)
270
        grad = self.fc1.backward(grad)
271

272
        grad = grad.reshape(-1, 12, 12, 16)
273
        grad = self.pool.backward(grad)
274

275
        grad = grad * relu_grad(self.after_relu2)
276
        grad = self.conv2.backward(grad)
277

278
        grad = grad * relu_grad(self.after_relu1)
279
        self.conv1.backward(grad)
280

281
    def step(self):
282
        if self.adam:
283
            self.adam.step([
284
                (self.fc2.W, self.fc2.dW),
285
                (self.fc2.b, self.fc2.db),
286
                (self.fc1.W, self.fc1.dW),
287
                (self.fc1.b, self.fc1.db),
288
                (self.conv2.W, self.conv2.dW),
289
                (self.conv2.b, self.conv2.db),
290
                (self.conv1.W, self.conv1.dW),
291
                (self.conv1.b, self.conv1.db),
292
            ])
293
        else:
294
            self.fc2.step(self.lr)
295
            self.fc1.step(self.lr)
296
            self.conv2.step(self.lr)
297
            self.conv1.step(self.lr)
298

299
    def predict_batch(self, X):
300
        """Return predicted classes for a batch."""
301
        y_pred = self.forward(X)
302
        return np.argmax(y_pred, axis=1)
303

304
    def __call__(self, x):
305
        """
306
        Single-sample predict wrapper.
307
        Accepts (28,28), (28,28,1) or (1,28,28,1) style inputs.
308
        """
309
        if x.ndim == 2:  # H, W
310
            x = x[None, :, :, None]
311
        elif x.ndim == 3:  # H, W, C
312
            x = x[None, :, :, :]
313
        return self.predict_batch(x)[0]
314

315
    def save(self, path):
316
        os.makedirs(os.path.dirname(path), exist_ok=True)
317
        np.savez(
318
            path,
319
            conv1_W=self.conv1.W,
320
            conv1_b=self.conv1.b,
321
            conv2_W=self.conv2.W,
322
            conv2_b=self.conv2.b,
323
            fc1_W=self.fc1.W,
324
            fc1_b=self.fc1.b,
325
            fc2_W=self.fc2.W,
326
            fc2_b=self.fc2.b,
327
        )
328

329
    def load(self, path):
330
        data = np.load(path)
331
        self.conv1.W = data["conv1_W"]
332
        self.conv1.b = data["conv1_b"]
333
        self.conv2.W = data["conv2_W"]
334
        self.conv2.b = data["conv2_b"]
335
        self.fc1.W = data["fc1_W"]
336
        self.fc1.b = data["fc1_b"]
337
        self.fc2.W = data["fc2_W"]
338
        self.fc2.b = data["fc2_b"]
339

340
# ============================================================
341
#                     Tiny Adam Optimizer
342
# ============================================================
343
class TinyAdam:
344
    """
345
    超简洁 Adam：以 (param, grad) 列表作为输入，直接原地更新参数。
346
    """
347
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
348
        self.lr = lr
349
        self.beta1 = beta1
350
        self.beta2 = beta2
351
        self.eps = eps
352
        self.t = 0
353
        self.m = {}
354
        self.v = {}
355

356
    def step(self, params_and_grads):
357
        self.t += 1
358
        for param, grad in params_and_grads:
359
            key = id(param)
360
            m = self.m.get(key, np.zeros_like(param))
361
            v = self.v.get(key, np.zeros_like(param))
362

363
            m = self.beta1 * m + (1 - self.beta1) * grad
364
            v = self.beta2 * v + (1 - self.beta2) * (grad ** 2)
365

366
            m_hat = m / (1 - self.beta1 ** self.t)
367
            v_hat = v / (1 - self.beta2 ** self.t)
368

369
            param -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
370

371
            self.m[key] = m
372
            self.v[key] = v
373

374
# ============================================================
375
#                         Train Helpers
376
# ============================================================
377
def train(model, X_train, y_train, epochs, batch_size):
378
    loss_list = []
379
    for epoch in range(epochs):
380
        pbar = tqdm.tqdm(
381
            range(0, len(X_train), batch_size),
382
            desc=f"Epoch {epoch+1}/{epochs}"
383
        )
384

385
        for i in pbar:
386
            X_batch = X_train[i:i+batch_size]
387
            y_batch = y_train[i:i+batch_size]
388

389
            y_pred = model.forward(X_batch)
390
            loss = cross_entropy(y_pred, y_batch)
391
            loss_list.append(loss)
392

393
            model.backward(y_batch)
394
            model.step()
395

396
            pbar.set_postfix({"loss": f"{loss:.4f}"})
397
    return loss_list
398

399
def test(model, X_test, y_test):
400
    y_pred_test = model.forward(X_test)
401
    pred = np.argmax(y_pred_test, axis=1)
402
    true = np.argmax(y_test, axis=1)
403
    acc = (pred == true).mean()
404
    print(f"\nTest accuracy: {acc:.4f}")
405
    return acc
406

407
def plot_loss(loss_list):
408
    plt.figure(figsize=(6,4))
409
    plt.plot(loss_list)
410
    plt.xlabel("Iteration")
411
    plt.ylabel("Loss (CE)")
412
    plt.yscale("log")
413
    plt.title("Minimal CNN Training Loss")
414
    plt.grid(True)
415
    plt.tight_layout()
416
    plt.show()
417
# ============================================================
418
#                         Train + Test
419
# ============================================================
420
if __name__ == "__main__":
421
    np.random.seed(0)
422

423
    # change to your dataset directory
424
    data_dir = "./mnist"
425

426
    # 1. Load MNIST (from your Kaggle files)
427
    X_train, y_train, X_test, y_test = load_mnist_from_local(data_dir)
428

429
    # 2. Create CNN model
430
    model = MinimalCNNClassifier(lr=0.001, use_adam=True)
431

432
    # 3. Train
433
    loss_list = train(model, X_train, y_train, epochs=3, batch_size=64)  # CNN 收敛快，3 epoch 就能 90%+
434

435
    # 4. Test accuracy
436
    acc = test(model, X_test, y_test)
437

438
    # 5. Plot loss
439
    plot_loss(loss_list)
440

441
    # 6. Save model
442
    save_path = os.path.join("models", "cnn_model.npz")
443
    model.save(save_path)
444
    print(f"Model saved to {save_path}!")
445

446
    # 7. Load model (optional)
447
    print("Reloading model...")
448
    model2 = MinimalCNNClassifier(lr=0.001, use_adam=True)
449
    model2.load(save_path)
450
    print("Reload complete.")
451

452
    # 8. Test reloaded model
453
    acc2 = test(model2, X_test, y_test)
454
    print(f"Reloaded model accuracy: {acc2:.4f}")
455

456
    # 9. Single-sample inference demo
457
    sample_pred = model2(X_test[0].squeeze())
458
    print(f"Single sample predicted class: {sample_pred}, true: {np.argmax(y_test[0])}")