Danh sách bài viết

Bài 16: nn.Module — xây dựng mạng neural đầu tiên

nn.Module là base class cho mọi neural network trong PyTorch: tự động track parameter cho autograd, quản lý nested submodule, save/load state_dict, move model sang GPU. Bài này đi qua cấu trúc subclass với __init__ và forward, nn.Linear, nn.Sequential, parameters/named_parameters, state_dict, train/eval mode, nn.ModuleList/ModuleDict, functional API, pitfall, và build một MLP cho MNIST.

24/05/2026
14 phút đọc
1 lượt xem
1

Mục tiêu bài học

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

  • Hiểu vai trò của nn.Module như base class — nơi tập trung parameter, submodule, state.
  • Viết được một network bằng cách subclass nn.Module với hai method __init__forward.
  • Dùng đúng nn.Linear, nn.Sequential, các layer phổ biến.
  • Truy cập model.parameters(), named_parameters(), đếm số tham số huấn luyện được.
  • Save / load checkpoint bằng state_dict, hiểu vì sao đó là cách khuyến nghị.
  • Switch train() / eval() đúng thời điểm — ảnh hưởng tới Dropout, BatchNorm.
  • Tránh các pitfall: quên super().__init__(), gán tensor thuần vào self, dùng Python list cho submodule.

Bài này nối tiếp Bài 15 — Autograd và là tiền đề cho Bài 17 — Optimizer SGD và Adam, nơi bạn sẽ ghép model.parameters() với optimizer để chạy training loop hoàn chỉnh.

2

nn.Module là gì

torch.nn.Module là base class trừu tượng cho mọi component trong neural network: từ một layer đơn (nn.Linear, nn.Conv2d) tới toàn bộ model (ResNet, BERT). Mỗi nn.Module giữ:

  • Parameter (nn.Parameter): tensor có requires_grad=True mà optimizer sẽ update.
  • Buffer: tensor đi cùng model nhưng không train (vd running_mean của BatchNorm).
  • Submodule: các nn.Module con — quản lý đệ quy.
  • Hook: callback chạy lúc forward/backward (debug, intermediate output).

Khi bạn gán một nn.Module hoặc nn.Parameter vào self.xxx trong __init__, PyTorch tự đăng ký vào internal dict (_modules, _parameters). Mọi thao tác sau đó — parameters(), to(device), state_dict() — đều dựa trên các dict này.

3

Vì sao dùng nn.Module

Về lý thuyết bạn có thể viết network bằng tensor và hàm thuần — nhưng nn.Module giải quyết bốn việc trở nên không tầm thường khi model lớn lên:

  • Tự động track parameter cho autograd: gọi model.parameters() ra đúng các tensor cần update; không cần bookkeeping bằng tay.
  • Save / load state_dict: serialize chỉ tham số, portable giữa các phiên bản code (mục 12).
  • Move toàn bộ model sang GPU dễ dàng: model.to("cuda") đệ quy đẩy mọi parameter và buffer xuống device — không cần loop tay từng tensor.
  • Quản lý nested submodule: encoder chứa nhiều block, mỗi block chứa attention + FFN... nn.Module đệ quy duyệt được toàn bộ cây.

Quan trọng nhất: contract giữa nn.Module và phần còn lại của PyTorch (optimizer, DataLoader, distributed wrapper, checkpoint, profiler) đều dựa vào interface này. Tự viết class ngoài nn.Module nghĩa là phải reimplement contract đó.

4

Cấu trúc subclass nn.Module

Khung tối thiểu cho một classifier MLP nhận ảnh MNIST flat 784 chiều, output 10 lớp:

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = MyModel()
x = torch.randn(32, 784)
logits = model(x)         # gọi model(x), KHÔNG gọi model.forward(x)
print(logits.shape)       # torch.Size([32, 10])

Hai điểm quan trọng:

  • super().__init__() phải gọi trước khi assign layer. Bỏ dòng này → các dict nội bộ chưa được khởi tạo → mọi assign sau đó không được track.
  • Gọi model(x) chứ không phải model.forward(x). nn.Module.__call__ chạy các hook trước/sau forward (cần cho register_forward_hook, autograd graph, profiler).
5

__init__ và forward

Hai method bắt buộc khi subclass nn.Module:

__init__: khởi tạo layer. Mỗi layer assign vào self.xxx được PyTorch ghi vào _modules. Đây là nơi quyết định kiến trúc — số chiều, số layer, type của từng layer.

forward: định nghĩa luồng dữ liệu từ input ra output. Đây là nơi quyết định thứ tự các op chạy, có thể chứa if, for, print Python (dynamic graph — nhắc lại B15).

Khi bạn gọi model(x):

  1. Python gọi model.__call__(x) (inherited từ nn.Module).
  2. __call__ chạy pre-forward hooks → gọi model.forward(x) → chạy forward hooks → trả output.

Không có method khác bị bắt buộc. Bạn có thể thêm method phụ trợ (extra_repr, reset_parameters, custom helper) nhưng chỉ __init__ + forward đủ chạy.

6

nn.Linear

nn.Linear(in_features, out_features, bias=True) implement một fully-connected layer:

\[ y = x W^\top + b \]

Trong đó tensor:

  • weight: shape (out_features, in_features).
  • bias: shape (out_features,); có thể tắt với bias=False.
layer = nn.Linear(784, 128)
print(layer.weight.shape)     # torch.Size([128, 784])
print(layer.bias.shape)       # torch.Size([128])

x = torch.randn(32, 784)      # batch 32
y = layer(x)
print(y.shape)                # torch.Size([32, 128])

Default init: weight theo Kaiming uniform với fan-in từ in_features, bias uniform trong khoảng nhỏ. Reset bằng layer.reset_parameters().

Lưu ý broadcast: nn.Linear chấp nhận input nhiều chiều miễn dim cuối khớp in_features — vd (batch, seq, in)(batch, seq, out) tự động. Đây là lý do nó dùng được trong Transformer block.

7

nn.Sequential

nn.Sequential là container chồng layer theo thứ tự. Forward = chạy từng layer liên tiếp.

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)
x = torch.randn(32, 784)
logits = model(x)             # (32, 10)

Đặt tên layer bằng OrderedDict:

from collections import OrderedDict

model = nn.Sequential(OrderedDict([
    ("fc1", nn.Linear(784, 128)),
    ("relu", nn.ReLU()),
    ("fc2", nn.Linear(128, 10)),
]))
print(model.fc1)              # Linear(in_features=784, out_features=128, bias=True)

Khi nào dùng Sequential:

  • Model linear-flow — input đi qua các layer theo đúng thứ tự khai báo, không nhánh.
  • Code ngắn, không cần subclass — thường dùng cho head của model lớn (classifier cuối ResNet) hoặc baseline.

Khi nào KHÔNG dùng được:

  • Có branching: skip connection (ResNet), Inception module — cần subclass.
  • Multiple input hoặc multiple output.
  • Control flow phụ thuộc input (if trong forward).
8

Common layers

Các layer hay gặp trong torch.nn:

  • Dense: nn.Linear.
  • Convolution: nn.Conv1d, nn.Conv2d, nn.Conv3d (Bài 28 sẽ đi sâu).
  • Activation: nn.ReLU, nn.Tanh, nn.Sigmoid, nn.Softmax, nn.GELU, nn.SiLU.
  • Normalization: nn.BatchNorm1d, nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm.
  • Regularization: nn.Dropout, nn.Dropout2d.
  • Pooling: nn.MaxPool2d, nn.AvgPool2d, nn.AdaptiveAvgPool2d.
  • Embedding: nn.Embedding(num_embeddings, embedding_dim) — lookup vector cho token ID.
  • Recurrent: nn.RNN, nn.LSTM, nn.GRU (Bài 34+).
  • Transformer: nn.TransformerEncoderLayer, nn.MultiheadAttention (Bài 36+).

Mỗi layer đều là nn.Module — có parameters(), state_dict(), to(device) như model lớn. Đây là lý do compose model lớn từ layer nhỏ đồng nhất.

9

model.parameters()

model.parameters() trả về generator các nn.Parameter — đệ quy qua mọi submodule:

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

for p in model.parameters():
    print(p.shape, p.requires_grad)
# torch.Size([128, 784]) True
# torch.Size([128])      True
# torch.Size([10, 128])  True
# torch.Size([10])       True

Use case chính: pass vào optimizer.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Optimizer giữ reference tới các parameter này; mỗi optimizer.step() update in-place dựa trên .grad đã có sẵn.

Lọc parameter cần freeze: set p.requires_grad = False trước khi pass vào optimizer, hoặc lọc trong list comprehension.

trainable = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(trainable, lr=1e-3)
10

model.named_parameters()

Khi cần debug hoặc apply per-layer learning rate, dùng named_parameters() để lấy kèm tên:

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = MyModel()
for name, p in model.named_parameters():
    print(name, p.shape)
# fc1.weight torch.Size([128, 784])
# fc1.bias   torch.Size([128])
# fc2.weight torch.Size([10, 128])
# fc2.bias   torch.Size([10])

Tên parameter = path tới nó trong cây module, nối bằng dấu chấm. Đây cũng là key trong state_dict (mục 12).

Ứng dụng phổ biến: weight decay khác nhau cho bias / norm vs weight:

decay, no_decay = [], []
for name, p in model.named_parameters():
    if "bias" in name or "norm" in name.lower():
        no_decay.append(p)
    else:
        decay.append(p)

optimizer = torch.optim.AdamW(
    [{"params": decay, "weight_decay": 0.01},
     {"params": no_decay, "weight_decay": 0.0}],
    lr=1e-3,
)
11

Đếm số parameter

Idiom phổ biến đếm tổng số tham số huấn luyện được:

num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable params: {num_params:,}")

Với MLP 784 → 128 → 10:

  • fc1.weight: 128 × 784 = 100,352.
  • fc1.bias: 128.
  • fc2.weight: 10 × 128 = 1,280.
  • fc2.bias: 10.
  • Tổng: 101,770.

Tham khảo mốc số tham số của các model thực tế: GPT-2 small ~124M, BERT-base ~110M, LLaMA-7B ~7B, GPT-4 (ước tính) hàng trăm B. Đếm parameter là check đầu tiên khi bạn đọc paper hoặc port model.

12

state_dict và load_state_dict

state_dict() trả về một OrderedDict ánh xạ tên parameter / buffer → tensor. Đây là cách PyTorch chuẩn để serialize model:

sd = model.state_dict()
print(list(sd.keys())[:4])
# ['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias']

# Save
torch.save(model.state_dict(), "model.pth")

# Load: tạo lại model cùng kiến trúc, rồi load weight
model2 = MyModel()
model2.load_state_dict(torch.load("model.pth", map_location="cpu"))
model2.eval()

Khuyến nghị: save state_dict thay vì cả object model. Lý do:

  • Pickle cả class có thể vỡ khi refactor code, đổi tên class, di chuyển file — phụ thuộc vào import path lúc load.
  • state_dict chỉ chứa tensor + tên, portable giữa các phiên bản code miễn kiến trúc khớp.
  • Dễ load partial (vd transfer learning: chỉ load encoder, bỏ classifier) qua strict=False:
missing, unexpected = model.load_state_dict(
    torch.load("pretrained.pth"), strict=False
)
print("missing keys:", missing)
print("unexpected keys:", unexpected)

Khi save checkpoint đầy đủ cho việc resume training, thường gói thêm optimizer state và epoch:

torch.save({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epoch": epoch,
}, "checkpoint.pth")

Từ PyTorch 2.4 trở đi, torch.load có flag weights_only=True để tránh deserialize arbitrary code — bật mặc định cho file không tin cậy.

13

model.to(device)

model.to(device) đẩy tất cả parameter và buffer của model (đệ quy qua submodule) sang device — in-place, return chính self để chain:

device = "cuda" if torch.cuda.is_available() else "cpu"
model = MyModel().to(device)

# Input cũng phải sang cùng device
x = torch.randn(32, 784, device=device)
logits = model(x)

Một số chi tiết hay quên:

  • Tensor thường (không phải parameter / buffer / submodule) gán vào self.foo sẽ không được di chuyển bởi to(device). Nếu cần đi cùng model, đăng ký bằng self.register_buffer("foo", tensor).
  • to(dtype) tương tự đổi dtype toàn model (vd model.to(torch.float16) cho inference half-precision).
  • Optimizer phải tạo sau khi model.to(device), vì optimizer giữ reference tới parameter tensor — nếu tạo trước, tensor bị move thành tensor mới ở device mới, reference cũ bị lệch.
14

model.train() và model.eval()

Hai method này switch model giữa train mode và eval mode — ảnh hưởng hành vi của một số layer:

  • nn.Dropout: train mode bật dropout; eval mode tắt (pass-through).
  • nn.BatchNorm*: train mode dùng mean/var của batch hiện tại và update running stats; eval mode dùng running stats đã tích luỹ.
  • nn.LayerNorm, nn.Linear, nn.Conv2d: hành vi không đổi giữa hai mode.
model.train()    # set self.training = True (đệ quy)
# ... train step ...

model.eval()     # set self.training = False (đệ quy)
with torch.no_grad():
    # ... eval step ...
    preds = model(x_val)

Pattern chuẩn cho training loop:

for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        # ... forward, backward, step ...
        pass

    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            # ... eval ...
            pass

Đây là một trong những lỗi im lặng phổ biến nhất khi mới làm DL: quên model.eval() trước khi đo validation → Dropout vẫn drop ngẫu nhiên, BatchNorm vẫn dùng batch stats → metric eval bị noise.

model.train() / eval() không liên quan đến torch.no_grad() — chúng giải quyết hai vấn đề khác nhau. Eval thường gồm cả hai: model.eval() để switch mode + torch.no_grad() để tắt autograd (B15).

15

Nested module

Model lớn thường tổ chức theo cây module. PyTorch tự discover submodule khi bạn assign một nn.Module vào self.xxx:

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 64)

    def forward(self, x):
        return torch.relu(self.fc2(torch.relu(self.fc1(x))))

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(64, 256)
        self.fc2 = nn.Linear(256, 784)

    def forward(self, z):
        return torch.sigmoid(self.fc2(torch.relu(self.fc1(z))))

class EncoderDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z)

model = EncoderDecoder()
for name, p in model.named_parameters():
    print(name)
# encoder.fc1.weight, encoder.fc1.bias, encoder.fc2.weight, ...
# decoder.fc1.weight, decoder.fc1.bias, ...

parameters(), state_dict(), to(device), train/eval đều đệ quy xuống encoder và decoder. Bạn không cần viết thêm dòng nào để wire chúng — đó là sức mạnh của contract nn.Module.

Khi debug, in print(model) ra cây kiến trúc readable, in model.encoder ra subtree riêng.

16

nn.ModuleList và nn.ModuleDict

Khi cần lưu một danh sách submodule (vd N block giống nhau trong Transformer), KHÔNG dùng Python list / dict thuần — submodule trong đó sẽ không được track. Dùng nn.ModuleListnn.ModuleDict:

class MLP(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(dims[i], dims[i + 1])
            for i in range(len(dims) - 1)
        ])

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers) - 1:
                x = torch.relu(x)
        return x

model = MLP([784, 256, 128, 64, 10])
print(sum(p.numel() for p in model.parameters()))   # 242,762

Đối chiếu với cách SAI:

class BrokenMLP(nn.Module):
    def __init__(self):
        super().__init__()
        # SAI: Python list — submodule không được register
        self.layers = [nn.Linear(784, 256), nn.Linear(256, 10)]

    def forward(self, x):
        return self.layers[1](torch.relu(self.layers[0](x)))

m = BrokenMLP()
print(list(m.parameters()))   # [] — rỗng!

nn.ModuleList hỗ trợ index, iterate, len(); nn.ModuleDict hỗ trợ truy cập bằng key, dùng khi có branching theo string (vd multi-task heads):

self.heads = nn.ModuleDict({
    "classification": nn.Linear(256, 10),
    "regression": nn.Linear(256, 1),
})
# forward: y_cls = self.heads["classification"](feat)

Lưu ý: nn.ModuleList không định nghĩa forward — bạn phải tự viết logic trong forward của module bao ngoài. Khác nn.Sequential ở chỗ này: Sequential tự chain còn ModuleList chỉ là container.

17

Functional API (torch.nn.functional)

Bên cạnh các Module có state (nn.Linear, nn.Conv2d...), PyTorch cung cấp phiên bản functional trong torch.nn.functional (quy ước import là F):

import torch.nn.functional as F

x = torch.randn(32, 10)
y1 = F.relu(x)
y2 = F.softmax(x, dim=-1)
y3 = F.cross_entropy(x, torch.randint(0, 10, (32,)))

Khác biệt:

  • Stateless: F.relu không có parameter, không cần khởi tạo. Gọi như hàm thường.
  • Module phiên bản (nn.ReLU, nn.Softmax) wrap hàm functional thành object — có thể assign vào self, xuất hiện trong print(model), tham gia state_dict (rỗng vì không state).

Khi nào dùng cái nào:

  • Activation không có parameter (ReLU, Tanh, Sigmoid, Softmax): cả hai đều đúng. Convention phổ biến: ưu tiên nn.X để đồng nhất với layer có parameter, và để print(model) hiển thị đầy đủ kiến trúc.
  • Layer có parameter (Linear, Conv, BatchNorm): luôn dùng nn.X trong __init__. Functional phiên bản (F.linear(x, w, b)) tồn tại nhưng bạn phải tự quản parameter — chỉ dùng khi viết custom layer.
  • Loss: F.cross_entropy(logits, target) tiện cho one-shot; nn.CrossEntropyLoss tiện khi cần config (weight, ignore_index).

Cùng một forward có thể viết hai cách tương đương:

class A(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)
        self.relu = nn.ReLU()
    def forward(self, x):
        return self.relu(self.fc(x))

class B(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)
    def forward(self, x):
        return F.relu(self.fc(x))

Cả hai tương đương về kết quả.

18

Common architectures

Cùng một nn.Module đủ để build mọi kiến trúc bạn sẽ gặp trong series này:

  • MLP cho tabular / vector input: stack nn.Linear + activation. Bài hiện tại minh hoạ với MNIST flat.
  • CNN cho image: nn.Conv2d + nn.BatchNorm2d + pooling + classifier head (Bài 27+).
  • RNN / LSTM / GRU cho sequence: nn.LSTM, nn.GRU wrap state recurrent (Bài 34+).
  • Transformer cho NLP / multi-modal: nn.MultiheadAttention, nn.TransformerEncoderLayer (Bài 36+).
  • Autoencoder / VAE / GAN: hai/ba module con (encoder, decoder, discriminator) chứa trong một module bao ngoài.

Đặc điểm chung: kiến trúc nào cũng là nn.Module. Khi đọc code framework lớn (HuggingFace Transformers, torchvision, timm), bạn sẽ thấy hàng nghìn class — tất cả đều subclass nn.Module theo cùng pattern bạn vừa học.

19

Pitfall thường gặp

  • Quên super().__init__(): AttributeError: cannot assign module before Module.__init__() call khi assign layer. Sửa: đặt super().__init__() ngay dòng đầu của __init__.
  • Assign tensor thuần vào self: self.mask = torch.ones(10) — tensor này không được track, không di chuyển khi to(device), không nằm trong state_dict. Sửa: self.register_buffer("mask", torch.ones(10)) nếu cần đi cùng model nhưng không train.
  • Dùng Python list / dict cho submodule: self.layers = [nn.Linear(...), ...] → parameter rỗng. Sửa: nn.ModuleList / nn.ModuleDict.
  • Gọi model.forward(x) trực tiếp: bỏ qua hook và một số guard nội bộ. Sửa: luôn gọi model(x).
  • Quên model.eval() trước validation → Dropout / BatchNorm sai mode → metric noisy hoặc bias.
  • Quên model.train() khi quay lại training sau eval → tương tự, nhưng nguy hiểm hơn vì model train với Dropout tắt.
  • Tạo optimizer trước model.to(device): optimizer giữ tensor ở device cũ. Sửa: to(device) trước, optimizer sau.
  • Save toàn bộ object thay vì state_dict: torch.save(model, ...) hoạt động nhưng dễ vỡ khi refactor.
  • Load state_dict không match kiến trúc: missing / unexpected keys → mặc định raise. Dùng strict=False chủ động khi transfer learning.
20

Code Python: MLP cho MNIST

Định nghĩa kiến trúc, đếm parameter, move sang GPU, save và load state_dict. Chưa có training loop — phần đó để Bài 17 (optimizer) và Bài 18 (DataLoader).

import torch
import torch.nn as nn
import torch.nn.functional as F

class MNISTClassifier(nn.Module):
    """MLP cho MNIST: 784 -> 128 -> 10."""

    def __init__(self, input_dim=784, hidden_dim=128, num_classes=10, p_drop=0.2):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.dropout = nn.Dropout(p_drop)
        self.fc2 = nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        # x shape: (batch, 1, 28, 28) hoặc (batch, 784)
        x = x.flatten(start_dim=1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)         # trả logits, không softmax (CrossEntropyLoss tự xử lý)


model = MNISTClassifier()
print(model)
# MNISTClassifier(
#   (fc1): Linear(in_features=784, out_features=128, bias=True)
#   (dropout): Dropout(p=0.2, inplace=False)
#   (fc2): Linear(in_features=128, out_features=10, bias=True)
# )

# Đếm parameter
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable params: {num_params:,}")   # 101,770

# Test forward
x = torch.randn(8, 1, 28, 28)
logits = model(x)
print(logits.shape)                           # torch.Size([8, 10])

# Save / load state_dict
torch.save(model.state_dict(), "mnist_mlp.pth")

model2 = MNISTClassifier()
model2.load_state_dict(torch.load("mnist_mlp.pth", map_location="cpu"))
model2.eval()

# So sánh output: giống hệt sau khi load
with torch.no_grad():
    out1 = model(x)
    out2 = model2(x)
print(torch.allclose(out1, out2))             # True (sau khi cả hai eval())

# Move sang GPU (nếu có)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
x = x.to(device)
print(model(x).device)                        # cuda:0 hoặc cpu

Một số quan sát:

  • Trong forward, x.flatten(start_dim=1) đổi (B, 1, 28, 28) thành (B, 784) — giữ batch dim, phẳng phần còn lại.
  • Output là logits, không phải xác suất. nn.CrossEntropyLoss hợp nhất log-softmax + NLL nên không cần softmax thủ công (đỡ instability).
  • Để output deterministic khi so sánh, cả hai model phải ở cùng mode (cả hai eval()) vì Dropout khác nhau giữa các forward khi training=True.
21

Bài tập

  1. Build một MLP 4-layer (784 → 256 → 128 → 64 → 10) bằng hai cách: (a) nn.Sequential, (b) subclass nn.Module. Verify hai version có cùng số parameter.
  2. Với MLP ở bài 1, in tên + shape của từng parameter bằng named_parameters(). Tính tay số parameter của từng layer rồi đối chiếu.
  3. Đếm tổng số parameter huấn luyện được của MLP 4-layer trên. So sánh với số parameter của nn.Linear(784, 10) đơn lẻ.
  4. Tạo model có 2 input (vd ảnh 784 chiều + vector metadata 16 chiều), concat trước layer cuối, output 10 lớp. Implement bằng subclass nn.Modulenn.Sequential không làm được vì có hai đường input.
  5. Save state_dict của model ở bài 1 ra "mlp.pth". Tạo một instance mới rồi load. Verify output trên cùng input giống nhau (sau khi cả hai eval()).
  6. Demo pitfall: tạo class subclass nn.Module nhưng dùng Python list cho danh sách layer. Verify model.parameters() ra rỗng. Sửa bằng nn.ModuleList, verify lần thứ hai.
  7. Tạo MLP có Dropout giữa các layer. So sánh output trên cùng input khi model.train() vs model.eval() — chạy nhiều lần để thấy variance của train mode.
  8. Cho một model, lọc named_parameters() thành hai nhóm: weight (áp weight decay) và bias + norm (không decay). In số parameter mỗi nhóm.
Đáp án ngắn
  1. Cùng số param vì cùng kiến trúc. Layer sizes: nn.Linear(784,256): 200,960; Linear(256,128): 32,896; Linear(128,64): 8,256; Linear(64,10): 650. Tổng = 242,762.
  2. Mỗi nn.Linear(a, b)weight shape (b, a) = \( a \cdot b \) param + bias shape (b,) = \( b \) param, tổng \( (a + 1) \cdot b \).
  3. 242,762 vs 7,850 — MLP nhiều hơn ~31 lần do hidden layer.
  4. Forward: concat = torch.cat([self.img_proj(img), self.meta_proj(meta)], dim=-1); out = self.classifier(concat).
  5. torch.allclose(model1(x), model2(x)) trả True sau khi cả hai eval().
  6. Pitfall: list(model.parameters()) trả []. Sửa self.layers = nn.ModuleList([...]) → ra đủ param.
  7. Train mode: output khác nhau giữa các lần forward (Dropout ngẫu nhiên). Eval mode: output deterministic.
  8. Group decay: tất cả *.weight trừ *norm*.weight. Group no_decay: *.bias + *norm*.weight. Thường tỉ lệ ~99:1.
22

Tóm tắt

  • nn.Module là base class cho mọi neural network — giữ parameter, buffer, submodule trong dict nội bộ.
  • Subclass cần hai method: __init__ khởi tạo layer (sau super().__init__()) và forward định nghĩa luồng dữ liệu.
  • Luôn gọi model(x), không gọi model.forward(x) trực tiếp.
  • nn.Linear(in, out) implement \( y = x W^\top + b \) với weight shape (out, in), bias shape (out,).
  • nn.Sequential chồng layer linear-flow; subclass khi cần branching, multi input/output, control flow.
  • parameters() / named_parameters() đệ quy qua submodule; dùng cho optimizer, debug, weight decay theo nhóm.
  • Save / load state_dict thay vì cả object model — portable hơn, dễ partial load với strict=False.
  • model.to(device) in-place đẩy mọi parameter và buffer sang device; tạo optimizer sau khi to(device).
  • model.train() / eval() switch mode — ảnh hưởng Dropout, BatchNorm. Bắt buộc trước eval / inference.
  • nn.ModuleList / nn.ModuleDict cho danh sách submodule, KHÔNG dùng Python list / dict thuần.
  • torch.nn.functional cung cấp phiên bản stateless; convention: ưu tiên nn.X cho consistency.
  • Pitfall hay gặp: quên super().__init__(), assign tensor thuần vào self, Python list cho submodule, quên eval() trước validation.