Kiểm tra shape và vector hóa tích chập bằng im2col

Ý tưởng gọn trong một câu: mỗi vị trí (i, j) của output tương ứng với đúng một patch K×K×C_in của input, nên ta trải mỗi patch thành một hàng. Ma trận kết quả có số hàng là _____ và mỗi hàng dài _____ phần tử.

Ví dụ tối giản, X 3×3, kernel 2×2, stride 1, không padding:

X = [[1,2,3],      cols =
     [4,5,6],        [1,2,4,5]
     [7,8,9]]        [2,3,5,6]
                     [4,5,7,8]
                     [5,6,8,9]

Bốn patch, bốn hàng. Lưu ý thứ tự hàng: duyệt n, rồi i, rồi j. Thứ tự này phải khớp chính xác với bước reshape ở dưới, nếu không kết quả sẽ đúng giá trị nhưng sai vị trí.

Cách dựng không cần vòng lặp trên từng pixel: lặp trên K*K vị trí kernel, mỗi lần cắt một slice có stride bằng S.

import numpy as np

def im2col(X, K, S=1, P=0):
    N, C, H, Wd = X.shape
    H_out = (H + 2*P - K)//S + 1
    W_out = (Wd + 2*P - K)//S + 1
    Xp = np.pad(X, ((0,0), (0,0), (P,P), (P,P)))
    cols = np.zeros((N, C, K, K, H_out, W_out))
    for ki in range(K):
        for kj in range(K):
            cols[:, :, ki, kj] = Xp[:, :, ki:ki+S*H_out:S, kj:kj+S*W_out:S]
    return cols.transpose(0, 4, 5, 1, 2, 3).reshape(N*H_out*W_out, C*K*K)

Hàm này là bản viết tay của torch.nn.functional.unfold(X, kernel_size=K, stride=S, padding=P), chỉ khác là unfold trả về shape (N, C*K*K, L) với L = H_out*W_out, tức bản transpose của cols.

Từ im2col đến một GEMM duy nhất

Trọng số (C_out, C_in, K, K) được reshape thành _____. Khi đó:

flowchart TD
    X["🖼️<br/>X (N,C,H,W)"] --> P["⬜<br/>Zero pad P"]
    P --> C["🧱<br/>im2col"]
    C --> M["📋<br/>cols (N·H_out·W_out, C·K·K)"]
    M --> G["✖️<br/>GEMM với W (C_out, C·K·K)"]
    G --> R["🔁<br/>Reshape + transpose"]
    R --> O["📤<br/>Out (N,C_out,H_out,W_out)"]

def conv_im2col(X, W, b, S=1, P=0):
    N, C, H, Wd = X.shape
    C_out, _, K, _ = W.shape
    H_out = (H + 2*P - K)//S + 1
    W_out = (Wd + 2*P - K)//S + 1

    cols = im2col(X, K, S, P)                      # (N*H_out*W_out, C*K*K)
    out = cols @ W.reshape(C_out, C*K*K).T         # (N*H_out*W_out, C_out)
    out = out.reshape(N, H_out, W_out, C_out)      # hàng đang xếp theo (n, i, j)
    out = out.transpose(0, 3, 1, 2)                # -> channel-first
    return out + b[None, :, None, None]

Hai dòng dễ sai nhất là dòng reshape và dòng transpose. Bạn phải transpose, vì reshape cho ra layout (N, H_out, W_out, C_out), còn mọi thứ trong khóa học (và cả naive conv của bài trước) dùng (N, C_out, H_out, W_out).

Bài tập 2. Bạn quên dòng transpose(0, 3, 1, 2). Chuyện gì xảy ra khi so với naive conv?

Xem đáp án

Không có exception nào, vì shape (N, H_out, W_out, C_out) vẫn hợp lệ. Giá trị thì sai bét: các trục bị hoán vị. Nếu tình cờ C_out ** H_out ** W_out thì shape trùng nhau hoàn toàn và lỗi còn im lặng hơn nữa. Đây chính là lý do phải luôn có một test đối chiếu giá trị, không chỉ kiểm tra shape.

Đối chiếu với naive conv

Vector hóa chỉ đáng tin khi có test so sánh. Dùng tensor nhỏ, float64, và so với bản bốn vòng lặp của bài trước:

rng = np.random.default_rng(0)
X = rng.standard_normal((2, 3, 9, 9))
W = rng.standard_normal((4, 3, 3, 3))
b = rng.standard_normal(4)

a = conv_naive(X, W, b, S=2, P=1)
c = conv_im2col(X, W, b, S=2, P=1)
print(a.shape, c.shape)          # (2, 4, 5, 5) (2, 4, 5, 5)
print(np.abs(a - c).max())       # ~1e-16
print(np.allclose(a, c, atol=1e-9))

Hai bẫy cần nhớ khi dùng slice có stride: ki:ki+S*H_out:S phải cắt đúng H_out phần tử, thiếu S*H_out là mất patch cuối; và np.pad phải pad cả hai trục không gian, không pad trục N hay C.

Cái giá của im2col: mỗi phần tử input bị sao chép khoảng _____ lần, nên bộ nhớ tăng theo hệ số đó. Với K=3 thì không đáng lo, với K=7 và batch lớn thì cols có thể to hơn cả ảnh gốc nhiều lần.

Bẫy shape thường gặp và mẹo kiểm tra nhanh

• Viết shape bằng tay trước, chạy code sau. Một dòng (8,3,32,32) -> (8,16,16,16) tiết kiệm hàng giờ debug.
• Đừng tin broadcasting. out + b[None, :, None, None] đúng, còn out + b sẽ cộng sai trục mà không báo lỗi nếu kích thước tình cờ khớp. Luôn in shape sau mỗi bước biến đổi.
• Thứ tự hàng của im2col là hợp đồng ngầm giữa im2col và reshape. Đổi một bên mà không đổi bên kia thì test giá trị sẽ bắt được, còn mắt thường thì không.
• Kiểm tra bằng float64 với tolerance 1e-9. Nếu bạn test bằng float32 rồi thấy sai lệch 1e-6, đó là nhiễu làm tròn chứ không phải bug logic.
• Ưu tiên tensor siêu nhỏ khi debug: (1, 2, 3, 3), kernel 2, stride 1. Ở kích thước đó bạn in được toàn bộ ma trận ra màn hình và tự kiểm tra bằng mắt.
• Quy tắc ngón tay: nếu H_out <= 0 thì cấu hình (K, S, P) không hợp lệ với input đó, kiểm tra ngay chứ đừng để numpy cấp phát mảng rỗng rồi lỗi ở chỗ khác.

Chốt lại và bước tiếp theo

Ba điều mang sang bài sau:

• Shape đầu ra tính bằng một công thức duy nhất, và phải khớp với tính tay trước khi chạy.
• im2col biến conv thành cols @ W.T, với cols shape (N*H_out*W_out, C*K*K) và W reshape thành (C_out, C*K*K).
• Mọi bước vector hóa cần một test np.allclose đối chiếu với naive conv trên tensor nhỏ.

Bài tiếp theo, lan truyền ngược qua lớp tích chập, sẽ dùng lại chính hai ma trận này: gradient theo W là một GEMM giữa cols.T và dout, còn gradient theo input là một phép "col2im" chạy ngược lại. Dựng đúng im2col hôm nay là bạn đã đi được nửa đường của backprop.

Front: Công thức H_out của lớp conv 2D? Back: H_out = floor((H + 2P - K)/S) + 1, và W_out tính tương tự trên W.

Front: Mỗi hàng của ma trận im2col chứa gì? Back: Một patch K×K×C_in đã trải phẳng, dài C_inKK phần tử.

Front: Số hàng của ma trận im2col? Back: N * H_out * W_out, xếp theo thứ tự (n, i, j).

Front: Shape của trọng số conv sau khi reshape để nhân với im2col? Back: (C_out, C_inKK).

Front: Kết quả cols @ W.T có shape gì trước khi reshape? Back: (NH_outW_out, C_out).

Front: Vì sao phải transpose sau reshape (N, H_out, W_out, C_out)? Back: Để về layout channel-first (N, C_out, H_out, W_out) mà naive conv và phần còn lại của mạng dùng.

Front: im2col đánh đổi điều gì? Back: Bộ nhớ tăng khoảng K*K lần để đổi lấy một phép GEMM duy nhất.

Front: Cách thêm zero padding đúng trong numpy? Back: np.pad(X, ((0,0),(0,0),(P,P),(P,P))), chỉ pad hai trục không gian.

Front: Slice nào lấy patch tại vị trí kernel (ki, kj)? Back: Xp[:, :, ki:ki+SH_out:S, kj:kj+SW_out:S].

Front: Cách kiểm tra im2col đúng là gì? Back: So np.allclose với naive conv trên tensor nhỏ, dùng float64, atol 1e-9.

Front: Với H=32, K=5, S=2, P=2 thì H_out bằng bao nhiêu? Back: 16.

Front: Bias được cộng vào đâu sau GEMM? Back: Cộng b[None, :, None, None] vào output, sau bước transpose.

Front: Hàm PyTorch tương đương im2col? Back: torch.nn.functional.unfold, trả về shape (N, CKK, L).