Project: forward pass và kiểm tra shape cho mạng 2 lớp

By Better Wise

Bạn đã có trong tay hai mảnh ghép: công thức softmax ổn định bằng log-sum-exp, và đạo hàm của cross-entropy theo logits. Bài này dựng nửa đầu của capstone: một forward pass chạy được cho mạng linear → ReLU → linear → softmax, với assert shape ở từng bước. Chưa có backward, chưa có training. Mục tiêu hẹp nhưng khắt khe: mọi tensor phải đúng shape mà bạn đã thiết kế trước khi viết dòng code đầu tiên, vì backward ở bài sau sẽ dựa hoàn toàn vào những shape này.

Hợp đồng shape: viết trước, code sau

Đây là thói quen của người làm research: chốt hợp đồng shape trên giấy rồi mới gõ code. Với N mẫu, D đặc trưng vào, H unit ẩn, C lớp:

Tensor	Shape	Ghi chú
`X`	`(N, D)`	mỗi hàng là một mẫu, đã flatten
`W1`, `b1`	`(D, H)`, `(H,)`	fan-in của `W1` là `D`
`Z1`	`(N, H)`	logits lớp ẩn
`A1`	`(N, H)`	ReLU elementwise, không đổi shape
`W2`, `b2`	`(H, C)`, `(C,)`	fan-in của `W2` là `H`
`Z2`	`(N, C)`	logits lớp ra, chưa softmax
`P`	`(N, C)`	xác suất, mỗi hàng tổng bằng 1

Nguyên tắc bất di bất dịch: bias luôn là vector 1D (H,) hoặc (C,), không phải (1, H). Broadcasting sẽ "cứu" bạn khi chạy, nhưng nó phá hợp đồng và làm assert nổ, đúng như mong muốn.

import numpy as np

def init_params(D, H, C, rng, scale=1e-2):
    return {
        "W1": rng.normal(0.0, scale, (D, H)),
        "b1": np.zeros(H),
        "W2": rng.normal(0.0, scale, (H, C)),
        "b2": np.zeros(C),
    }

def forward(X, p):
    N, D = X.shape
    W1, b1, W2, b2 = p["W1"], p["b1"], p["W2"], p["b2"]
    assert W1.shape[0] == D, "fan-in của W1 phải bằng số đặc trưng"

    Z1 = X @ W1 + b1
    assert Z1.shape == (N, W1.shape[1])

    A1 = np.maximum(Z1, 0.0)
    assert A1.shape == Z1.shape, "ReLU không được đổi shape"

    Z2 = A1 @ W2 + b2
    assert Z2.shape == (N, b2.shape[0])
    assert np.isfinite(Z2).all(), "logits có NaN/inf, kiểm tra scale khởi tạo"

    return {"Z1": Z1, "A1": A1, "Z2": Z2, "N": N}

Lưu ý: forward trả về cache chứ không chỉ Z2. Bài sau cần Z1, A1, Z2 để tính gradient, nên dựng thói quen lưu ngay từ bây giờ.

flowchart TD
    X["🖼️<br/>X (N, D)"] --> L1["➕<br/>Z1 = X·W1 + b1<br/>(N, H)"]
    L1 --> R["⚡<br/>A1 = ReLU(Z1)<br/>(N, H)"]
    R --> L2["➕<br/>Z2 = A1·W2 + b2<br/>(N, C)"]
    L2 --> S["🎯<br/>P = softmax(Z2)<br/>(N, C)"]

Và lớp softmax, viết tách riêng để test độc lập:

def softmax(Z):
    assert Z.ndim == 2
    Zs = Z - Z.max(axis=1, keepdims=True)   # (N, 1) để broadcast
    E = np.exp(Zs)
    P = E / E.sum(axis=1, keepdims=True)
    assert P.shape == Z.shape
    assert np.allclose(P.sum(axis=1), 1.0, atol=1e-6)
    return P

Important: keepdims=True không phải chi tiết vụn. Bỏ nó, Z.max(axis=1) trả về (N,), trừ vào (N, C) sẽ broadcast sai trục và bạn nhận về một ma trận sai lệch mà không có lỗi nào được ném ra. Đây là loại bug tệ nhất: chạy được, kết quả sai.

Bài tập 1: chốt con số trước khi code

Với D = 784, H = 64, C = 10: tổng số tham số của mạng là bao nhiêu? Nếu batch N = 8 thì tensor lớn nhất trong forward có bao nhiêu phần tử?

Đáp án: tham số $= 784 \cdot 64 + 64 + 64 \cdot 10 + 10 = 50{.}176 + 64 + 640 + 10 = 50{.}890$. Tensor lớn nhất là Z1 và A1 với $8 \cdot 64 = 512$ phần tử. Con số nhỏ này là lý do mạng 2 lớp là nơi tốt nhất để debug bằng tay.

Lộ trình: sáu mốc, mỗi mốc một assert

Sinh dữ liệu giả lập: X = rng.normal(size=(N, D)), y = rng.integers(0, C, size=N). Assert X.shape == (N, D) và y.shape == (N,), dtype nguyên.
Viết init_params và assert W1.shape == (D, H), W2.shape == (H, C), bias là 1D.
Dựng lớp tuyến tính 1, assert Z1.shape == (N, H).
Thêm ReLU, assert A1.shape == Z1.shape và (A1 >= 0).all().
Dựng lớp tuyến tính 2, assert Z2.shape == (N, C) và np.isfinite(Z2).all().
Bọc softmax, assert P.shape == (N, C) và mỗi hàng tổng bằng 1 trong sai số 1e-6.

Chạy mốc 6 với ba bộ tham số biên: N = 1, C = 2, và H = 1. Mạng nhỏ đến mức bạn nhìn được từng phần tử thì bug shape không có chỗ trốn.

Bài tập 2: dữ liệu ảnh chưa flatten

Bạn nạp được X shape (32, 28, 28) và p["W1"] shape (784, 64). Dòng X @ W1 sẽ ném lỗi gì, và sửa thế nào cho đúng hợp đồng?

Gợi ý: @ chỉ nhận ma trận 2D. Sửa bằng X = X.reshape(X.shape[0], -1) trước khi gọi forward, rồi assert X.shape == (32, 784). Đừng reshape bên trong lớp tuyến tính, vì lớp đó không nên biết dữ liệu gốc là ảnh hay bảng.

Dấu hiệu bạn đã xong

forward chạy không lỗi với (N, D) = (8, 784), (1, 784), và (8, 5).
Tất cả assert shape vẫn bật khi chạy "thật", không bị comment out.
P.sum(axis=1) xấp xỉ 1 với atol=1e-6, và P không có NaN.
forward trả về dict có đủ Z1, A1, Z2 để bài sau dùng.
Đổi H từ 64 sang 128 thì code vẫn chạy, chỉ có shape trong assert đổi theo. Không có hằng số ma thuật nào hard-code 64.

Bẫy shape và vài gợi ý cứu hộ

Bẫy	Triệu chứng	Cách xử lý
`b1` khởi tạo `(1, H)`	Chạy được, assert sai	Luôn `np.zeros(H)`
Softmax trước khi lưu logits	Backward bài sau mất đầu vào	Trả `Z2`, tính `P` riêng
`scale` khởi tạo quá lớn	logits lớn, softmax bão hòa	Dùng `1e-2` hoặc He theo fan-in
Quên flatten ảnh	lỗi `matmul` ở `X @ W1`	`reshape(N, -1)` ở tầng dữ liệu
Dùng `np.max(A1)` thay vì `np.maximum`	ReLU "co" cả batch về một số	`np.maximum(Z1, 0.0)`

Một gợi ý về quy trình: viết một hàm check_shapes(p, X) in ra toàn bộ shape một lần, gọi ở đầu forward khi còn debug. Mười giây đọc output rẻ hơn nhiều so với mười phút đoán mò.

Bài tập 3: kiểm tra softmax có thật sự ổn định

Cho Z = np.array([[1000.0, 0.0], [0.0, 0.0]]). Nếu bạn lấy np.exp(Z) trực tiếp rồi chia, chuyện gì xảy ra? Bản có trừ max thì sao?

Đáp án: bản không trừ max sinh inf / inf và trả về NaN ở hàng đầu. Bản trừ max đưa hàng đầu về [0, -1000], exp cho [1, 0], kết quả [1, 0]. Cùng một xác suất, khác nhau ở chỗ một bên là NaN. Đây chính là lý do assert np.isfinite(P).all() đáng được thêm vào.

Nghiệm thu xong thì mang gì sang bài sau

Bạn vừa dựng xong bộ khung: hợp đồng shape bằng bảng, init_params, lớp tuyến tính, ReLU, lớp tuyến tính thứ hai, softmax ổn định, và một cache đầy đủ Z1, A1, Z2. Mọi thứ đều đã được kiểm bằng assert nên khi sang bài kế tiếp, bug shape sẽ không còn là nghi phạm.

Bài tiếp theo mở đúng cái cache này ra: viết backward pass, kiểm chứng bằng numerical gradient, và ghép thành training loop cho đến khi loss giảm.

Front: Với X shape (N, D) và W1 shape (D, H), Z1 có shape gì? Back: (N, H).

Front: ReLU có làm đổi shape của tensor không? Back: Không, nó elementwise nên A1 cùng shape với Z1.

Front: Shape của W2 và b2 khi có C lớp? Back: W2 là (H, C), b2 là (C,).

Front: Vì sao phải trừ max theo từng hàng trước khi lấy exp trong softmax? Back: Để tránh overflow: exp của số lớn là inf, trừ max giữ mọi logit ≤ 0 nên exp ≤ 1.

Front: keepdims=True trong Z.max(axis=1, keepdims=True) để làm gì? Back: Giữ shape (N, 1) để phép trừ broadcast đúng lên (N, C).

Front: Sau softmax, tổng mỗi hàng bằng bao nhiêu? Back: Xấp xỉ 1, kiểm bằng np.allclose(P.sum(axis=1), 1.0).

Front: Forward pass cần cache những gì cho backward? Back: Z1, A1 và Z2 (P nếu dùng để tính loss).

Front: Vì sao assert A1.shape == Z1.shape đáng giá? Back: Nó bắt lỗi dùng nhầm phép làm đổi shape, ví dụ np.max với axis sai.

Front: Với D=784, H=64, C=10, tổng số tham số là bao nhiêu? Back: 784·64 + 64 + 64·10 + 10 = 50.890.

Front: X shape (32, 28, 28) muốn đưa vào lớp tuyến tính thì làm gì? Back: reshape thành (32, 784), tức (N, -1).

Front: Vì sao bias nên là (H,) thay vì (1, H)? Back: Broadcasting vẫn chạy nhưng phá hợp đồng shape, làm assert sai và gây lệch ở backward.

Front: Khởi tạo W với scale quá lớn gây ra chuyện gì? Back: Logits lớn, softmax bão hòa, loss cao và gradient gần 0; dùng scale nhỏ như 1e-2 hoặc He.

Front: Logits trong mạng này là tensor nào? Back: Z2, đầu ra tuyến tính trước softmax, shape (N, C).

Với X shape (N, D) và W1 shape (D, H), logits của lớp ẩn có shape _____. ReLU là phép toán _____ nên không làm đổi shape của tensor. Ma trận trọng số W2 phải có shape _____, với C là số lớp. Để softmax ổn định, ta trừ đi _____ trước khi lấy exp. Sau softmax, tổng mỗi hàng của P phải xấp xỉ _____. Nếu X có shape (32, 28, 28) thì phải _____ trước khi đưa vào lớp tuyến tính.