Bài 30: Flatten và Fully Connected ở cuối CNN

Sau bài học, bạn sẽ:

• Hiểu vai trò của Flatten + FC ở cuối CNN: chuyển feature map 2D sang vector và ánh xạ sang số class.
• Đọc được kiến trúc LeNet, AlexNet ở mức "Conv stack + FC head".
• Biết cách tính output size sau mỗi Conv/Pool và suy ngược ra in_features của FC đầu tiên.
• Phân biệt Flatten+FC truyền thống và Global Average Pooling (GAP) modern; biết khi nào dùng cái nào.
• Dùng đúng nn.Flatten, nn.Linear, nn.AdaptiveAvgPool2d, nn.LazyLinear.
• Tránh các bug thường gặp: sai in_features, quên Flatten, double softmax với CrossEntropyLoss.

Bài này nối tiếp Bài 29 — Pooling Layer (kết quả của backbone là feature map nhỏ) và chuẩn bị cho Bài 31 — CNN cho CIFAR-10 (ráp Conv + Pool + Flatten + FC thành mô hình hoàn chỉnh).

Sau khi đi qua một dãy Conv + Pool, dữ liệu vẫn ở dạng tensor 3 chiều có shape \( (C, H, W) \) — chiều rộng / cao đã nhỏ lại (vd \( 7 \times 7 \) hoặc \( 6 \times 6 \)), số kênh đã tăng (vd 64, 128, 256). Mỗi vị trí spatial là một "feature vector" mô tả đặc trưng cục bộ của một patch ảnh.

Đầu ra của bài phân loại lại là một vector class probability: với 10 class thì cần shape \( (10,) \), với 1000 class thì \( (1000,) \). Không thể trực tiếp ánh xạ tensor \( (C, H, W) \) sang \( (\text{num\_classes},) \) bằng Conv — cần module có thể "trộn" mọi vị trí spatial lại với nhau.

Hai lựa chọn phổ biến:

• Flatten + FC (kinh điển): duỗi \( (C, H, W) \) thành vector dài \( C \cdot H \cdot W \), rồi đưa qua một hoặc vài nn.Linear để giảm chiều và ra số class.
• Global Average Pooling + Linear (modern): lấy trung bình mỗi kênh trên toàn bộ spatial dim, ra vector \( (C,) \), rồi một nn.Linear(C, \text{num\_classes}) duy nhất.

Một CNN classification cổ điển luôn chia làm hai phần rõ ràng:

• Feature extractor / backbone: stack của Conv + Activation + Pool, có nhiệm vụ trích đặc trưng từ low-level (edge, texture) đến high-level (part, object). Input là ảnh \( (3, H, W) \), output là feature map \( (C, h, w) \) với \( h, w \) đã nhỏ đi nhiều lần.
• Classifier head: bộ phận biến feature map sang vector logit. Truyền thống là Flatten → Linear → ReLU → Dropout → Linear → ... → Linear(num_classes).

Cách phân tách này quan trọng vì:

• Backbone có thể tái sử dụng qua các task khác nhau (classification, detection, segmentation) — cơ sở của transfer learning (bài 32).
• Khi đếm số parameter, classifier head của CNN cổ điển chiếm đa số (vd AlexNet: ~60M tổng thì ~58M ở FC head).

Flatten chỉ là một thao tác reshape, không có parameter. Nó duỗi tensor multi-dim thành vector:

• Single sample: \( (C, H, W) \rightarrow (C \cdot H \cdot W,) \).
• Batch: \( (B, C, H, W) \rightarrow (B, C \cdot H \cdot W) \).

Trong PyTorch có hai cách dùng phổ biến:

import torch
import torch.nn as nn

x = torch.randn(8, 64, 7, 7)   # batch=8, C=64, H=W=7

# Cách 1: module nn.Flatten — mặc định flatten từ dim=1 (giữ batch dim)
flat = nn.Flatten()
print(flat(x).shape)           # torch.Size([8, 3136])

# Cách 2: method tensor — chỉ rõ start_dim
print(x.flatten(1).shape)      # torch.Size([8, 3136])

Số \( 3136 = 64 \cdot 7 \cdot 7 \). Giữ dim 0 vì đó là chiều batch — KHÔNG được flatten luôn batch dim, nếu không một batch sẽ "trộn" với nhau.

view(-1, C * H * W) hoặc reshape cũng cho kết quả tương tự nhưng phải tự tính số. nn.Flatten rõ ràng và an toàn hơn trong nn.Sequential.

LeNet-5 (LeCun et al. 1998) là CNN classification thực dụng đầu tiên, dùng cho nhận dạng số viết tay. Cấu trúc với input MNIST đã pad lên \( 32 \times 32 \):

Backbone đẩy output về \( (16, 5, 5) \). Flatten cho 400 feature. Ba lớp FC \( 400 \to 120 \to 84 \to 10 \). Tổng param ~62 nghìn, trong đó FC chiếm 95%.

AlexNet (Krizhevsky et al. 2012, arXiv:1102.0183 reformulation) thắng ILSVRC-2012, kích hoạt làn sóng deep learning. Cấu trúc dạng "Conv stack + FC head":

• Input: \( (3, 224, 224) \).
• 5 lớp Conv + 3 MaxPool xen kẽ, output cuối backbone: \( (256, 6, 6) \).
• Flatten → 9216 feature.
• FC 4096 → ReLU → Dropout.
• FC 4096 → ReLU → Dropout.
• FC 1000 (số class ImageNet).

Tổng ~60M parameter, trong đó:

• FC1 (9216 → 4096): \( 9216 \cdot 4096 \approx 37.7\text{M} \).
• FC2 (4096 → 4096): \( 4096 \cdot 4096 \approx 16.8\text{M} \).
• FC3 (4096 → 1000): \( 4096 \cdot 1000 \approx 4.1\text{M} \).

Riêng FC head ~58.6M, gấp đôi-ba lần param của toàn bộ Conv. Đây là động lực khiến các kiến trúc sau (Network-in-Network, ResNet) đi tìm cách giảm chi phí phần FC — dẫn đến GAP (mục 10).

Một CNN nhỏ cho MNIST (\( 28 \times 28 \)) với 2 Conv + 2 Pool + Flatten + FC:

import torch.nn as nn

model = nn.Sequential(
    # ----- Feature extractor / backbone -----
    nn.Conv2d(1, 32, kernel_size=3, padding=1),   # (32, 28, 28)
    nn.ReLU(),
    nn.MaxPool2d(2),                              # (32, 14, 14)

    nn.Conv2d(32, 64, kernel_size=3, padding=1),  # (64, 14, 14)
    nn.ReLU(),
    nn.MaxPool2d(2),                              # (64, 7, 7)

    # ----- Classifier head -----
    nn.Flatten(),                                 # (64 * 7 * 7 = 3136,)
    nn.Linear(64 * 7 * 7, 128),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(128, 10),                           # logit cho 10 class
)

Conv dùng padding=1, kernel \( 3 \times 3 \), stride mặc định 1 nên giữ nguyên spatial size. MaxPool \( 2 \times 2 \) stride 2 chia đôi. Sau 2 lần Pool: \( 28 / 4 = 7 \). Số kênh cuối là 64 nên flatten ra \( 64 \cdot 7 \cdot 7 = 3136 \) — đúng số phải truyền vào nn.Linear.

Để biết in_features của FC đầu tiên, cần track output size qua từng layer. Công thức từ bài 28–29:

\[ H_{\text{out}} = \left\lfloor \frac{H_{\text{in}} + 2p - k}{s} \right\rfloor + 1 \]

Một vài case thường dùng:

• Conv2d(kernel=3, padding=1, stride=1): \( H_{\text{out}} = H_{\text{in}} \). Giữ nguyên size.
• Conv2d(kernel=5, padding=2, stride=1): cũng giữ nguyên size (padding "same" với kernel lẻ là \( (k-1)/2 \)).
• MaxPool2d(2, stride=2): \( H_{\text{out}} = H_{\text{in}} / 2 \) (làm tròn xuống).
• Conv2d(kernel=3, padding=0, stride=1): \( H_{\text{out}} = H_{\text{in}} - 2 \).

Áp dụng ngược: với MNIST \( 28 \times 28 \), sau 2 lần Pool \( 2 \times 2 \): \( 28 \to 14 \to 7 \). Cuối backbone \( 64 \cdot 7 \cdot 7 = 3136 \) feature. Nếu input là \( 32 \times 32 \) (vd CIFAR), cùng backbone sẽ ra \( 64 \cdot 8 \cdot 8 = 4096 \) — phải sửa in_features của FC tương ứng.

Khi input có nhiều size khác nhau (vd ảnh người dùng upload), kích thước cuối backbone cũng thay đổi → in_features của FC sai. nn.AdaptiveAvgPool2d(target_size) ép output về size cố định bất kể input:

pool = nn.AdaptiveAvgPool2d((7, 7))   # luôn ra (C, 7, 7)
print(pool(torch.randn(1, 64, 13, 11)).shape)  # (1, 64, 7, 7)
print(pool(torch.randn(1, 64, 50, 40)).shape)  # (1, 64, 7, 7)

Trong nội bộ, PyTorch tính kernel + stride sao cho output đúng target_size. Pattern phổ biến của các kiến trúc modern:

nn.AdaptiveAvgPool2d((1, 1)),   # (C, 1, 1)
nn.Flatten(),                   # (C,)
nn.Linear(C, num_classes),

Đây chính là Global Average Pooling (mục 10) đóng gói thành module.

Global Average Pooling (GAP) lấy trung bình mỗi kênh trên toàn bộ chiều spatial:

\[ \text{GAP}(X)_c = \frac{1}{H \cdot W} \sum_{i=1}^{H} \sum_{j=1}^{W} X_{c, i, j} \]

Output: \( (C, H, W) \rightarrow (C, 1, 1) \rightarrow \text{Flatten} \rightarrow (C,) \). Sau đó dùng MỘT nn.Linear(C, \text{num\_classes}) duy nhất.

GAP do Lin et al. (2013, arXiv:1312.4400, "Network in Network") đề xuất, sau đó được ResNet (He et al. 2015, arXiv:1512.03385), EfficientNet, MobileNet... áp dụng làm chuẩn. Trong PyTorch:

classifier_gap = nn.Sequential(
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),
    nn.Linear(64, 10),
)

So với Flatten+FC, GAP có ưu điểm:

• Ít param: không cần 9216×4096 weight; chỉ \( C \cdot \text{num\_classes} \) ở Linear cuối.
• Robust với input size khác nhau: GAP luôn ra \( (C,) \) bất kể \( H, W \).
• Regularize tự nhiên: ép mỗi feature map "đại diện" cho một concept; không có FC để "ghi nhớ" pattern cụ thể của training data — giảm overfitting.

ResNet-50 đạt top-1 ~76% ImageNet với chỉ ~25M param, trong khi VGG-16 cần ~138M để được ~71%. Phần lớn cách biệt đến từ classifier head: GAP loại bỏ hàng trăm triệu param và giữ accuracy nhờ backbone sâu hơn.

Khi nào vẫn dùng Flatten+FC:

• Bài rất nhỏ (MNIST, FashionMNIST với input \( 28 \times 28 \)) — backbone không deep, GAP có thể quá agressive.
• Khi spatial size cuối backbone không quá lớn (vd \( 7 \times 7 \)) và bạn muốn FC học mixing tự do.
• Mục đích pedagogy — học CNN với LeNet style.

FC head với hàng triệu param dễ overfitting. Dropout (bài 22) là kỹ thuật regularization phổ biến đặt giữa các FC layer:

nn.Flatten(),
nn.Linear(3136, 256),
nn.ReLU(),
nn.Dropout(0.5),       # tắt ngẫu nhiên 50% neuron khi train
nn.Linear(256, 10),

AlexNet đặt Dropout \( p = 0.5 \) sau hai FC 4096 đầu tiên — bài báo gốc gọi đây là yếu tố quan trọng giúp tránh overfitting trên ImageNet. VGG cũng theo pattern tương tự.

Khi dùng GAP + 1 FC duy nhất, Dropout ít cần thiết hơn — bản thân GAP đã có hiệu ứng regularize. ResNet, EfficientNet không dùng Dropout trong classifier head (chỉ dùng Stochastic Depth ở residual block).

Một lỗi thường gặp khi chuyển từ tutorial sang code thật là apply Softmax ở layer cuối khi loss đã có sẵn:

• Multi-class classification: dùng nn.CrossEntropyLoss. Loss này đã kết hợp LogSoftmax + NLLLoss. KHÔNG được thêm nn.Softmax ở output — sẽ thành "double softmax", gradient kém và loss nhỏ giả tạo.
• Binary classification: dùng nn.BCEWithLogitsLoss. Đã kết hợp Sigmoid + BCELoss ở dạng numerically stable. KHÔNG thêm nn.Sigmoid.

Layer cuối của model là nn.Linear(..., num_classes) — output là logit, không phải xác suất. Chỉ khi inference mới apply softmax để lấy xác suất hiển thị:

model.eval()
with torch.no_grad():
    logits = model(x)
    probs = torch.softmax(logits, dim=1)
    pred = probs.argmax(dim=1)

Khi training, chỉ cần đẩy logit thẳng vào loss(logits, target).

Lưu ý multi-label khác multi-class: ảnh có thể vừa có "mèo" vừa có "chó" — mỗi nhãn là một bài binary độc lập, không loại trừ nhau. Vì vậy dùng sigmoid per-label thay vì softmax.

Khi lắp xong model, luôn check shape của tensor qua từng layer trước khi train. Cách đơn giản với nn.Sequential:

import torch

x = torch.randn(1, 1, 28, 28)
for layer in model:
    x = layer(x)
    print(f"{layer.__class__.__name__:<12} -> {tuple(x.shape)}")

Output (ví dụ):

Conv2d       -> (1, 32, 28, 28)
ReLU         -> (1, 32, 28, 28)
MaxPool2d    -> (1, 32, 14, 14)
Conv2d       -> (1, 64, 14, 14)
ReLU         -> (1, 64, 14, 14)
MaxPool2d    -> (1, 64, 7, 7)
Flatten      -> (1, 3136)
Linear       -> (1, 128)
ReLU         -> (1, 128)
Dropout      -> (1, 128)
Linear       -> (1, 10)

Khi mô hình phức tạp hơn (custom nn.Module với forward), có thể dùng torchinfo:

from torchinfo import summary
summary(model, input_size=(1, 1, 28, 28))

• Sai in_features của FC đầu tiên: lỗi RuntimeError: mat1 and mat2 shapes cannot be multiplied (B x A and B' x C). Fix: in lại shape qua từng layer (mục 15) để biết số đúng.
• Quên Flatten: đưa thẳng tensor \( (B, C, H, W) \) vào nn.Linear sẽ raise lỗi shape vì Linear kỳ vọng \( (..., \text{in\_features}) \).
• Apply Softmax + CrossEntropyLoss: double softmax. Loss vẫn tính được nhưng gradient sai, accuracy thấp hơn nhiều so với kỳ vọng. Bỏ nn.Softmax ở model.
• Flatten luôn batch dim: dùng x.flatten() không tham số sẽ duỗi cả batch dim → một batch bị trộn. Dùng x.flatten(1) hoặc nn.Flatten().
• Đổi input size mà quên đổi in_features: train MNIST OK rồi chuyển CIFAR (\( 32 \times 32 \)) là crash. Hoặc dùng AdaptiveAvgPool2d để loại bỏ ràng buộc, hoặc dùng LazyLinear (mục 17).

Từ PyTorch 1.8, nn.LazyLinear(out_features) không yêu cầu chỉ định in_features — nó sẽ tự suy ra từ shape của tensor lần forward đầu tiên:

model = nn.Sequential(
    nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Flatten(),
    nn.LazyLinear(128),    # in_features tự suy từ forward đầu tiên
    nn.ReLU(),
    nn.Linear(128, 10),
)

# Forward một dummy tensor để initialize:
_ = model(torch.randn(1, 1, 28, 28))
print(model[-3])   # Linear(in_features=3136, out_features=128, bias=True)

Tiện khi prototype hoặc khi muốn chuyển giữa nhiều input size. Lưu ý:

• PyTorch raise warning rằng API này còn experimental — interface có thể thay đổi.
• Phải forward ít nhất một tensor trước khi save state_dict, nếu không weight chưa được tạo.
• Trong code production, vẫn nên ghi rõ in_features sau khi đã xác định size.

Hai phần của CNN có vai trò khác nhau về mặt functional:

• Backbone (Conv + Pool): học feature từ low-level (edge, texture) đến high-level (part, object). Feature này tương đối chung — backbone đã pretrain trên ImageNet vẫn hữu ích cho ảnh y khoa, ảnh vệ tinh...
• Classifier head (FC): ánh xạ feature sang class của task cụ thể. Rất task-specific.

Đây là nền tảng của transfer learning (bài 32): lấy backbone đã pretrain, thay classifier head bằng head mới cho task của mình, rồi:

• Frozen backbone: chỉ train head. Dùng khi data nhỏ, giống ImageNet domain (ảnh tự nhiên).
• Fine-tune toàn bộ: train cả backbone với LR nhỏ hơn. Dùng khi data đủ lớn hoặc domain khác xa (ảnh y khoa).

Trong torchvision, lấy ResNet và thay head chỉ là:

import torchvision.models as models

resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
# resnet.fc là Linear(2048, 1000) cho ImageNet
resnet.fc = nn.Linear(2048, num_classes)   # thay head

Backbone của ResNet kết thúc bằng GAP, output \( (2048,) \) — chỉ cần một nn.Linear mới là xong.

So sánh hai phiên bản cho MNIST, chung backbone, khác classifier head:

import torch
import torch.nn as nn

def make_backbone():
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    )

# Variant A: Flatten + FC
model_a = nn.Sequential(
    make_backbone(),
    nn.Flatten(),
    nn.Linear(64 * 7 * 7, 128),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(128, 10),
)

# Variant B: GAP + FC
model_b = nn.Sequential(
    make_backbone(),
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),
    nn.Linear(64, 10),
)

def count_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"Flatten+FC : {count_params(model_a):,} param")
print(f"GAP+FC     : {count_params(model_b):,} param")

x = torch.randn(1, 1, 28, 28)
print("Out A:", model_a(x).shape)   # torch.Size([1, 10])
print("Out B:", model_b(x).shape)   # torch.Size([1, 10])

Kết quả tiêu biểu:

Flatten+FC : 421,962 param
GAP+FC     : 19,530 param

Cùng backbone, GAP+FC giảm hơn 20 lần param classifier head. Trên MNIST cả hai đạt accuracy gần nhau (~99%) — nhưng trên ImageNet sự khác biệt mới rõ: backbone deep + GAP học feature tốt hơn nhiều so với việc dùng FC khổng lồ.

1. Build một CNN cho MNIST có 2 Conv + 2 Pool + Flatten + FC. Viết tay công thức tính output size sau từng layer, so sánh với output thực khi forward một tensor \( (1, 1, 28, 28) \).
2. Sửa lại model trên cho CIFAR-10 (\( 32 \times 32 \), 3 kênh). Cập nhật in_features của FC đầu tiên cho đúng.
3. Viết lại version GAP+FC: thay phần classifier bằng AdaptiveAvgPool2d((1, 1)) + Flatten + Linear(channels, 10). Print count_params cả hai version, ghi lại tỉ lệ.
4. Thay AdaptiveAvgPool2d((1, 1)) bằng AdaptiveAvgPool2d((4, 4)) + Flatten + Linear. Tính số param mới, giải thích sự khác biệt.
5. Thử apply nn.Softmax(dim=1) ở output rồi train với CrossEntropyLoss. So sánh accuracy với version không apply softmax — quan sát hiệu ứng double softmax.
6. Build một model với nn.LazyLinear thay vì Linear(64*7*7, 128). Forward một dummy tensor, in ra in_features được tự suy.

• CNN classification = backbone (Conv + Pool) + classifier head (Flatten + FC, hoặc GAP + FC).
• Flatten duỗi \( (B, C, H, W) \) thành \( (B, C \cdot H \cdot W) \) — không có parameter; dùng nn.Flatten hoặc x.flatten(1).
• FC đầu tiên cần in_features = C * H * W sau backbone — tính qua công thức Conv/Pool hoặc inspect shape bằng dummy forward.
• LeNet và AlexNet dùng Flatten + nhiều FC; ResNet, EfficientNet, MobileNet dùng GAP + 1 FC duy nhất.
• GAP giảm số param classifier head 20–100 lần, robust với input size khác nhau, có hiệu ứng regularize.
• AdaptiveAvgPool2d(target_size) ép output về size cố định; AdaptiveAvgPool2d((1, 1)) chính là GAP.
• Layer cuối là Linear(..., num_classes) → output là logit. Không apply Softmax/Sigmoid khi dùng CrossEntropyLoss/BCEWithLogitsLoss.
• Dropout giữa các FC giúp tránh overfitting (AlexNet, VGG); GAP-based head ít cần Dropout.
• nn.LazyLinear tự suy in_features — tiện khi prototype.
• Sự tách biệt backbone / head là nền tảng của transfer learning (bài 32).