Bài 19: Dataset và DataLoader — quản lý batch dữ liệu

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

• Hiểu vì sao PyTorch tách Dataset và DataLoader thay vì viết vòng for trực tiếp.
• Subclass torch.utils.data.Dataset đúng cách: override __len__ và __getitem__.
• Dùng TensorDataset khi data đã ở dạng tensor.
• Cấu hình DataLoader với batch_size, shuffle, num_workers, pin_memory, drop_last.
• Viết collate_fn để gom sample variable-length thành batch (padding).
• Load MNIST bằng torchvision.datasets với pipeline transforms.
• Phân chia tập train / val bằng random_split; biết Sampler dùng để làm gì.
• Tránh các bug hay gặp: num_workers trên Windows / Notebook, quên transform cho test set.

Bài 18 (Training Loop) đã giả định có sẵn loader để lặp qua batch. Bài này giải thích chính cái loader đó được dựng ra như thế nào, và làm sao đẩy data từ file / RAM / network thành tensor batch sạch.

Training một model deep learning luôn cần ba việc lặp đi lặp lại với dữ liệu:

• Chia batch: model train trên mini-batch (32, 64, 128…), không phải toàn bộ tập một lượt.
• Shuffle mỗi epoch: phá thứ tự để gradient không bias theo cấu trúc dataset.
• Load song song: trong khi GPU đang tính step i, CPU đã đọc + tiền xử lý sample cho step i+1.

Tự viết các việc này bằng vòng for tay được, nhưng tốn code và dễ bug (off-by-one ở batch cuối, thiếu shuffle, không tận dụng được nhiều worker). PyTorch chuẩn hoá thành hai abstraction:

• Dataset — tả dữ liệu: có bao nhiêu sample, lấy sample thứ i ra như thế nào.
• DataLoader — chạy qua dữ liệu: gom thành batch, shuffle, dùng nhiều worker.

Tách hai vai trò giúp đổi dataset không phải sửa training loop, và ngược lại đổi batch size / num_workers không phải sửa code dataset.

torch.utils.data.Dataset là một abstract class — bạn kế thừa và override hai method:

• __len__(self) — trả về tổng số sample. DataLoader cần biết để chia batch.
• __getitem__(self, idx) — nhận một int index, trả về một sample. Quy ước phổ biến: trả tuple (X, y) cho supervised.

PyTorch còn có khái niệm IterableDataset (cho stream không biết trước độ dài) — ít gặp ở mức beginner; bài này chỉ dùng Dataset kiểu map-style (có __len__).

Một sample không bắt buộc là tuple. Bạn có thể trả dict {"input": X, "label": y, "id": ...}, hữu ích khi có nhiều input. Tuple là quy ước mặc định mà nhiều helper code giả định.

Khuôn mẫu chuẩn khi data đã ở dạng array / tensor trong RAM:

import torch
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

Cách dùng:

import numpy as np

X = np.random.randn(1000, 20).astype(np.float32)
y = np.random.randint(0, 3, size=1000)

ds = MyDataset(X, y)
print(len(ds))         # 1000
x0, y0 = ds[0]
print(x0.shape, y0)    # torch.Size([20]) tensor(0)

• Convert sang tensor trong __init__ giúp tránh chi phí convert lặp lại mỗi lần lấy sample.
• dtype phải đúng ngay từ đây: feature thường float32; label classification phải int64 (long) cho CrossEntropyLoss.
• Khi dataset rất nhỏ và đơn giản như trên, TensorDataset (mục 5) là cách ngắn hơn.

TensorDataset là Dataset dựng sẵn cho trường hợp mọi thứ đã là tensor:

from torch.utils.data import TensorDataset

X_tensor = torch.randn(1000, 20)
y_tensor = torch.randint(0, 3, (1000,))

dataset = TensorDataset(X_tensor, y_tensor)
print(dataset[0])
# (tensor([...]), tensor(label))

• Nhận bao nhiêu tensor tuỳ ý, miễn là chiều 0 trùng nhau. __getitem__(i) trả tuple (t1[i], t2[i], ...).
• Dùng khi feature engineering xong (đã tensor hoá) và muốn nhanh.
• Không hỗ trợ transform / augmentation trên-the-fly; muốn vậy phải subclass thủ công.

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=False,
)

for X_batch, y_batch in loader:
    # X_batch shape: (32, 20) - hoặc nhỏ hơn ở batch cuối nếu drop_last=False
    # y_batch shape: (32,)
    ...

Tính chất quan trọng:

• DataLoader là iterable, không phải iterator: mỗi lần for batch in loader bắt đầu một epoch mới (shuffle lại nếu bật).
• Số batch một epoch = len(loader) = ceil(len(dataset) / batch_size) (hoặc floor nếu drop_last=True).
• Default collate_fn sẽ stack các sample (vốn là tuple tensor) thành tuple batch tensor — phù hợp với phần lớn dữ liệu fixed-shape.

• batch_size: thường 16, 32, 64, 128, 256. Lớn hơn → train nhanh hơn nếu RAM/VRAM đủ, nhưng quá lớn có thể làm xấu generalization (Series 3 Module 3 sẽ bàn). Giá trị mặc định 1.
• shuffle=True cho train, False cho validation và test. Khi shuffle=True, mỗi epoch random permutation các index — bạn không cần tự gọi shuffle.
• num_workers: số process song song load data. 0 = load trong main process (chậm nhưng dễ debug). Trên Linux server, 4-8 là khởi điểm phổ biến; cân theo số CPU core và tốc độ disk.
• pin_memory=True: cấp phát batch trong page-locked memory ở CPU; sau đó .to(device, non_blocking=True) sang GPU nhanh hơn. Chỉ ích khi train trên GPU.
• drop_last=True: nếu len(dataset) không chia hết batch_size, bỏ batch cuối (nhỏ hơn). Cần khi dùng BatchNorm với batch rất nhỏ (1 sample) gây lỗi; hoặc khi muốn shape batch cố định cho code downstream.
• persistent_workers=True: giữ worker process sống qua các epoch, tránh chi phí spawn lại. Cần khi __init__ của dataset nặng (load file index lớn).
• prefetch_factor: mỗi worker prefetch trước bao nhiêu batch. Default 2; tăng nếu pipeline I/O chậm.

Lưu ý num_workers không phải càng lớn càng nhanh: vượt số core hoặc băng thông disk thì bottleneck chuyển sang chỗ khác. Đo bằng time qua một epoch là cách thực dụng nhất.

Sau khi DataLoader gọi dataset[i] cho batch_size index, nó cần ghép các sample thành tensor batch. Việc đó do collate_fn đảm nhận. Mặc định: torch.utils.data.default_collate sẽ torch.stack các tensor cùng shape — vừa đủ cho ảnh đã resize hoặc feature fixed-length.

Khi sample có shape khác nhau (text độ dài khác, graph số node khác), default collate sẽ lỗi. Lúc đó viết custom collate_fn:

from torch.nn.utils.rnn import pad_sequence
import torch

def custom_collate(batch):
    # batch = list of (X, y), X có shape (T_i, D), y scalar
    Xs = [item[0] for item in batch]
    ys = torch.tensor([item[1] for item in batch])

    # padding về cùng T_max, batch_first=True -> (B, T_max, D)
    Xs_padded = pad_sequence(Xs, batch_first=True, padding_value=0.0)
    lengths = torch.tensor([x.shape[0] for x in Xs])
    return Xs_padded, lengths, ys

loader = DataLoader(dataset, batch_size=32, collate_fn=custom_collate)

• collate_fn chỉ nhận batch (list các sample) và trả về cái mà bạn muốn nhận trong vòng for.
• Trả thêm lengths hữu ích khi dùng pack_padded_sequence (Module 5: RNN) hoặc tạo attention mask.
• Custom collate cũng dùng để tokenize on-the-fly cho text (nhưng phải nhớ làm trong worker process — đắt CPU).

torchvision.datasets đóng gói các dataset ảnh phổ biến cùng logic download và parse:

• MNIST, FashionMNIST — chữ số / quần áo 28×28 grayscale.
• CIFAR10, CIFAR100 — ảnh màu 32×32.
• ImageNet — chỉ wrapper, dataset phải tự tải.
• ImageFolder — chuyển một thư mục có cấu trúc root/class_x/*.jpg thành dataset có label.

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])

train_set = datasets.MNIST(
    root="./data",
    train=True,
    download=True,
    transform=transform,
)
test_set = datasets.MNIST(
    root="./data",
    train=False,
    download=True,
    transform=transform,
)

print(len(train_set), len(test_set))    # 60000 10000
img, label = train_set[0]
print(img.shape, label)                 # torch.Size([1, 28, 28]) 5

0.1307 và 0.3081 là mean / std đã tính sẵn cho MNIST. Với dataset khác, tự tính trên tập train hoặc tra docs.

torchvision.transforms là chuỗi các bước biến đổi ảnh; mỗi bước nhận và trả ảnh (PIL.Image hoặc tensor):

• Resize((H, W)), CenterCrop(size) — đổi kích thước.
• RandomCrop(size, padding=4), RandomHorizontalFlip(p=0.5), RandomRotation(degrees=15) — data augmentation. Chỉ dùng cho tập train; tập val/test phải deterministic.
• ToTensor() — convert PIL.Image hoặc numpy.ndarray (H, W, C) uint8 [0,255] thành tensor (C, H, W) float [0, 1].
• Normalize(mean, std) — áp dụng per-channel: (x - mean) / std. Phải gọi sau ToTensor().
• Compose([...]) chain nhiều bước theo thứ tự.

Ví dụ điển hình cho CIFAR-10 train:

train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2470, 0.2435, 0.2616)),
])

eval_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2470, 0.2435, 0.2616)),
])

Lưu ý RandomCrop, RandomHorizontalFlip không có trong eval transform — đây là quy tắc bất di bất dịch: augmentation chỉ áp dụng cho train. Tập eval phải lặp lại được, không random.

Augmentation chi tiết hơn (cutout, mixup, autoaugment) sẽ ở Bài 26.

Bên cạnh torch.utils.data.Dataset, thư viện HuggingFace datasets cung cấp hàng ngàn dataset NLP / multimodal với API riêng:

from datasets import load_dataset

ds = load_dataset("imdb")            # text classification
print(ds["train"][0])
# {'text': 'I rented I AM CURIOUS-YELLOW ...', 'label': 0}

Để dùng với PyTorch:

• ds.set_format("torch") để các cột tensor trả về dưới dạng torch.Tensor.
• ds["train"] hành xử như list-like Dataset map-style — có __len__ và __getitem__, dùng thẳng được trong DataLoader sau khi tokenize.
• Mạnh ở phần mapping (ds.map(tokenize_fn, batched=True)) và caching (kết quả lưu Arrow file trên đĩa, đọc lần sau cực nhanh).

Bài này không đi sâu vào tokenizer — sẽ gặp lại trong Series 4 (LLM & GenAI). Tạm hiểu: HuggingFace datasets là một format dataset song song, dùng cho NLP tiện hơn đi Dataset tay.

Hai chiến lược load khi data trên đĩa:

• In-memory: __init__ load toàn bộ vào RAM. __getitem__ chỉ index. Đơn giản, nhanh khi đã warm-up; chỉ khả thi nếu data đủ nhỏ.
• Lazy: __init__ chỉ lưu danh sách path / index. __getitem__ mở và đọc 1 file khi cần. Khả thi với dataset hàng triệu ảnh; tốc độ phụ thuộc disk và num_workers.

Khuôn mẫu lazy cho dataset ảnh:

from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset

class ImageFolderDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = Path(root)
        self.files = sorted(self.root.glob("*/*.jpg"))   # ./root/class/img.jpg
        self.classes = sorted({p.parent.name for p in self.files})
        self.cls_to_idx = {c: i for i, c in enumerate(self.classes)}
        self.transform = transform

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        path = self.files[idx]
        img = Image.open(path).convert("RGB")
        label = self.cls_to_idx[path.parent.name]
        if self.transform is not None:
            img = self.transform(img)
        return img, label

torchvision.datasets.ImageFolder chính là phiên bản dùng được luôn của khuôn mẫu này.

Có một dataset to, muốn chia train / val:

from torch.utils.data import random_split

full = MyDataset(X, y)
n_train = int(0.8 * len(full))
n_val = len(full) - n_train

train_ds, val_ds = random_split(
    full, [n_train, n_val],
    generator=torch.Generator().manual_seed(42),
)

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=64, shuffle=False)

• random_split trả về Subset (wrap dataset gốc + list index) — vẫn dùng được như Dataset bình thường.
• Truyền generator với seed để chia lặp lại được giữa các lần chạy.
• Lưu ý: random_split không hỗ trợ stratify theo class — nếu dataset imbalanced và cần giữ tỉ lệ, dùng sklearn.model_selection.train_test_split trên index rồi tạo Subset thủ công.

Subset(dataset, indices) cho phép tự chỉ định các index muốn lấy — hữu ích khi muốn lấy con tập theo điều kiện (ví dụ chỉ class 0 và 1).

Khi cần thứ tự sample khác mặc định, dùng sampler:

• SequentialSampler — duyệt theo index 0, 1, 2… (mặc định khi shuffle=False).
• RandomSampler — random permutation (mặc định khi shuffle=True).
• WeightedRandomSampler — sample theo weight cho từng index. Hữu ích khi class imbalanced: cho class hiếm weight cao hơn để model nhìn thấy đều. Chi tiết về imbalanced ở Series 2 Bài 42 / Series 3 Bài 26.
• SubsetRandomSampler — random trong tập index cho trước.
• BatchSampler — gom các index thành batch trước khi qua __getitem__; dùng khi cần kiểm soát thành phần mỗi batch.

Ví dụ WeightedRandomSampler cho 3 class với tỉ lệ 70/20/10:

from torch.utils.data import WeightedRandomSampler

class_count = [700, 200, 100]
weights_per_class = [1.0 / c for c in class_count]
# weight cho mỗi sample = 1 / số sample trong class của nó
sample_weights = [weights_per_class[label] for _, label in dataset]

sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(dataset),
    replacement=True,
)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
# KHÔNG truyền shuffle=True khi đã có sampler

Khi truyền sampler, không được set shuffle=True cùng lúc (DataLoader sẽ raise).

Để hai lần chạy cho cùng thứ tự batch (debug, so sánh model):

def seed_worker(worker_id):
    import numpy as np, random
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(42)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    worker_init_fn=seed_worker,
    generator=g,
)

• generator=g kiểm soát thứ tự shuffle trong main process.
• worker_init_fn seed cho từng worker process — quan trọng khi code transform hoặc augmentation dùng numpy / random bên trong (mỗi worker được fork với cùng seed → nếu không reseed, mọi worker random giống hệt nhau!).
• Còn cần torch.manual_seed, torch.cuda.manual_seed_all để model init và dropout deterministic (xem Bài 13 mục 15).

• num_workers > 0 trên Windows / Notebook: PyTorch dùng multiprocessing theo phương thức spawn trên Windows, mọi code top-level sẽ chạy lại trong worker → nếu không gói trong if __name__ == "__main__": sẽ vô hạn spawn. Trong Jupyter Notebook tương tự, thường buộc phải num_workers=0.
• Quên transform cho test: train normalize với mean/std nhất định, test không normalize → input phân bố khác hẳn → metric tệ một cách khó hiểu. Luôn dùng cùng Normalize cho train và eval.
• Augmentation trên test: ngược lại, đem RandomCrop / RandomFlip sang eval transform → kết quả không lặp lại được, đánh giá lệch.
• DataLoader trả tuple thay vì dict: nếu Dataset.__getitem__ trả dict, default collate cũng trả dict batch (mỗi key là batch tensor). Hay nhầm là tuple. Kiểm tra bằng next(iter(loader)) trước khi nhúng vào loop.
• Shuffle val/test: gây mismatch khi so kết quả batch-by-batch giữa các run. Luôn shuffle=False cho eval.
• pin_memory=True trên CPU-only: vô hại nhưng không tăng tốc — bật chỉ khi train trên GPU.
• Sửa self.X in-place trong __getitem__: do worker được fork với copy-on-write, sửa data trong worker có thể không phản chiếu về main, nhưng cũng dễ gây bug khó lần. Luôn trả tensor mới hoặc copy trước khi sửa.

Pipeline 1 — Custom Dataset từ NumPy + DataLoader:

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

class TabularDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

X = np.random.randn(500, 10).astype(np.float32)
y = np.random.randint(0, 3, size=500)

ds = TabularDataset(X, y)
loader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=0)

xb, yb = next(iter(loader))
print(xb.shape, yb.shape, xb.dtype, yb.dtype)
# torch.Size([64, 10]) torch.Size([64]) torch.float32 torch.int64

Pipeline 2 — MNIST với torchvision + transforms:

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])

train_set = datasets.MNIST("./data", train=True, download=True, transform=transform)
test_set  = datasets.MNIST("./data", train=False, download=True, transform=transform)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True,
                          num_workers=2, pin_memory=True)
test_loader  = DataLoader(test_set,  batch_size=256, shuffle=False,
                          num_workers=2, pin_memory=True)

for xb, yb in train_loader:
    print(xb.shape, yb.shape)   # torch.Size([128, 1, 28, 28]) torch.Size([128])
    break

Pipeline 3 — Variable-length sequence với custom collate:

import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

class SeqDataset(Dataset):
    def __init__(self, samples):
        # samples = list of (seq_tensor shape (T_i,), label int)
        self.samples = samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]

def collate(batch):
    seqs = [item[0] for item in batch]
    labels = torch.tensor([item[1] for item in batch])
    lengths = torch.tensor([s.shape[0] for s in seqs])
    seqs_padded = pad_sequence(seqs, batch_first=True, padding_value=0)
    return seqs_padded, lengths, labels

samples = [
    (torch.arange(3),  0),
    (torch.arange(7),  1),
    (torch.arange(5),  0),
    (torch.arange(2),  1),
]

loader = DataLoader(SeqDataset(samples), batch_size=4, collate_fn=collate)
seqs, lengths, labels = next(iter(loader))
print(seqs.shape)    # torch.Size([4, 7])  - T_max = 7
print(lengths)       # tensor([3, 7, 5, 2])
print(labels)        # tensor([0, 1, 0, 1])

1. Subclass Dataset nhận một list các tuple (X_i, y_i) nơi X_i là tensor 1D, y_i là int. Override __len__, __getitem__. Tạo 100 sample fake và in len(ds), ds[0].
2. Wrap dataset trên bằng DataLoader với batch_size=16, shuffle=True. Lấy next(iter(loader)) và in shape của X_batch, y_batch.
3. Dùng torchvision.datasets.MNIST với transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]). In len tập train, shape của 1 sample.
4. Viết collate_fn nhận list [(seq_tensor, label), ...] với seq_tensor độ dài khác nhau, dùng pad_sequence với padding_value=-1, trả về (padded, lengths, labels). Test với 3 sample độ dài 2, 5, 4.
5. Lấy train_set MNIST, dùng random_split chia 80/20 thành train và val; seed = 42. In len mỗi phần.
6. Tạo TensorDataset với X = torch.randn(500, 4), y = torch.randint(0, 3, (500,)). Tạo WeightedRandomSampler để mỗi class có cùng tần suất kỳ vọng trong batch (mặc dù tỉ lệ class trong dataset không cân). Iterate 10 batch và đếm tần suất mỗi class.

• PyTorch tách hai vai trò: Dataset mô tả "có gì, lấy thế nào"; DataLoader lo batching, shuffle, parallel load.
• Subclass Dataset bằng cách override __len__ và __getitem__; TensorDataset là shortcut khi data đã là tensor.
• Tham số DataLoader cần nhớ: batch_size, shuffle, num_workers, pin_memory, drop_last; thêm persistent_workers, prefetch_factor khi tối ưu.
• collate_fn custom dùng cho variable-length (text, graph); thường kèm pad_sequence và trả thêm lengths.
• torchvision.datasets + transforms phủ phần lớn nhu cầu ảnh; ToTensor() rồi Normalize(mean, std); augmentation chỉ trên tập train.
• HuggingFace datasets là format song song, mạnh cho NLP và caching.
• Lazy load (đọc file trong __getitem__) cho dataset lớn; in-memory khi nhỏ.
• random_split chia train/val; Subset chọn index tuỳ ý; Sampler kiểm soát thứ tự (đặc biệt WeightedRandomSampler cho imbalanced).
• Reproducibility cần seed cả generator của DataLoader và worker_init_fn cho từng worker.
• Bug hay gặp: num_workers > 0 trên Windows/Notebook, quên transform cho test, augmentation đem sang eval, shuffle=True ở val/test.

Bài 20 sẽ quay về model: cách khởi tạo trọng số (Xavier, He) ảnh hưởng thế nào đến hội tụ, và vì sao mặc định của nn.Linear đã hợp lý cho phần lớn trường hợp.