Pytorch基礎教程：Dataset與DataLoader加載數據實戰

數據加載的基石：PyTorch中的Dataset與DataLoader¶

在機器學習和深度學習的訓練流程中，數據加載是至關重要的一環。如果數據加載環節混亂或低效，會直接影響模型訓練的效率和穩定性。PyTorch提供了Dataset和DataLoader兩個核心組件，幫助我們高效地管理和加載數據。本文將通過通俗易懂的語言和實戰代碼，帶初學者掌握數據加載的基本流程。

一、爲什麼需要Dataset和DataLoader？¶

在訓練模型時，我們需要將數據按批次（batch）輸入模型。Dataset負責“存儲和讀取單個數據樣本”，而DataLoader負責“將Dataset中的數據打包成批次，並支持打亂順序、多線程加載等優化”。沒有它們，我們需要手動處理數據的索引、打亂順序、批量讀取等複雜操作，容易出錯且效率低下。

二、Dataset：數據的“倉庫”¶

Dataset是PyTorch中處理數據的抽象基類，我們需要繼承它並實現兩個核心方法：
- __getitem__(self, index)：返回索引爲index的單個數據樣本（通常包含特徵和標籤）。
- __len__(self)：返回數據集的總樣本數。

1. 簡單自定義Dataset¶

假設我們有一組簡單的二維數據，每個樣本包含20個特徵和1個標籤（0-4的整數）。我們可以自定義一個Dataset類：

import torch
from torch.utils.data import Dataset

# 自定義Dataset類
class SimpleDataset(Dataset):
    def __init__(self, features, labels):
        # features: 所有樣本的特徵張量，形狀爲 (樣本數, 特徵數)
        # labels: 所有樣本的標籤張量，形狀爲 (樣本數,)
        self.features = features
        self.labels = labels

    def __getitem__(self, index):
        # 返回單個樣本的特徵和標籤
        return self.features[index], self.labels[index]

    def __len__(self):
        # 返回數據集總樣本數
        return len(self.features)

# 生成隨機數據
features = torch.randn(1000, 20)  # 1000個樣本，每個20個特徵
labels = torch.randint(0, 5, (1000,))  # 1000個標籤，範圍0-4

# 創建Dataset實例
dataset = SimpleDataset(features, labels)

# 測試：獲取第0個樣本
sample_features, sample_labels = dataset[0]
print(f"第0個樣本特徵形狀: {sample_features.shape}, 標籤: {sample_labels}")

2. 使用內置子類TensorDataset（更簡單）¶

如果數據已經是張量形式，TensorDataset可以直接包裝特徵和標籤，無需自定義類：

from torch.utils.data import TensorDataset

# 同樣生成隨機數據（與上文一致）
features = torch.randn(1000, 20)
labels = torch.randint(0, 5, (1000,))

# 直接用TensorDataset包裝
dataset = TensorDataset(features, labels)  # 自動處理__getitem__和__len__

# 測試：獲取第1個樣本
sample_features, sample_labels = dataset[1]
print(f"第1個樣本特徵形狀: {sample_features.shape}, 標籤: {sample_labels}")

三、DataLoader：數據的“快遞員”¶

DataLoader是Dataset的“包裝器”，它會從Dataset中按批次（batch）讀取數據，並提供以下功能：
- batch_size：每批數據的樣本數（必填）。
- shuffle：訓練時是否打亂數據順序（訓練集設爲True，驗證/測試集設爲False）。
- num_workers：多線程加載數據的數量（需注意Windows系統下默認設爲0）。
- pin_memory：是否鎖定內存（優化GPU傳輸數據速度，可選）。

1. DataLoader基礎用法¶

from torch.utils.data import DataLoader

# 創建DataLoader
dataloader = DataLoader(
    dataset,                # 輸入Dataset
    batch_size=32,          # 每批32個樣本
    shuffle=True,           # 訓練時打亂順序
    num_workers=0,          # 單線程加載（Windows系統建議設爲0，避免報錯）
    pin_memory=True         # 優化GPU數據傳輸（可選）
)

# 迭代DataLoader，獲取批次數據
for batch_idx, (batch_features, batch_labels) in enumerate(dataloader):
    print(f"批次 {batch_idx} | 特徵形狀: {batch_features.shape}, 標籤形狀: {batch_labels.shape}")
    # 例如：批次0 | 特徵形狀: torch.Size([32, 20]), 標籤形狀: torch.Size([32])
    break  # 只打印第一個批次

2. 關鍵參數解釋¶

batch_size：每批數據的大小（如64、128）。如果總樣本數不是batch_size的整數倍，最後一批會自動截斷（或用drop_last=True丟棄最後不完整的批次）。
shuffle：訓練時必須設爲True，避免模型“學習到數據順序”（如MNIST的數字順序），驗證/測試時設爲False保證結果可復現。
num_workers：多線程加載數據（數值越大，加載速度越快，但需考慮內存佔用）。Windows系統默認num_workers=0，Linux/macOS可設爲CPU核心數（如os.cpu_count()）。

四、實戰：加載圖像數據（以MNIST爲例）¶

除了上述張量數據，我們經常需要加載圖像數據（如MNIST手寫數字）。PyTorch的torchvision庫提供了內置數據集和數據轉換工具，結合Dataset和DataLoader可快速實現。

1. 安裝依賴（如果未安裝）¶

pip install torch torchvision

2. 加載MNIST數據集¶

import torchvision
import torchvision.transforms as transforms

# 定義數據轉換：將圖像轉爲張量並歸一化
transform = transforms.Compose([
    transforms.ToTensor(),       # 圖像轉張量（0-1範圍）
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST數據集的均值和標準差
])

# 加載訓練集和測試集
train_dataset = torchvision.datasets.MNIST(
    root='./data',             # 數據存儲路徑
    train=True,                # 加載訓練集
    download=True,             # 首次運行自動下載數據
    transform=transform        # 應用數據轉換
)

test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,               # 加載測試集
    download=True,
    transform=transform
)

# 創建DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=2              # Linux/macOS可設爲2（示例）
)

test_loader = DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False,
    num_workers=2
)

# 測試：迭代查看圖像和標籤
for images, labels in train_loader:
    print(f"圖像形狀: {images.shape} (batch_size=64, 通道=1, 高=28, 寬=28)")
    print(f"標籤形狀: {labels.shape}")
    break

五、常見問題與解決¶

錯誤：Windows下num_workers>0報錯
解決：Windows系統下num_workers必須設爲0（或根據PyTorch版本調整，建議設爲0簡化代碼）。
數據形狀不符合預期
解決：打印batch_features.shape確認維度是否正確。例如圖像數據通常爲(batch_size, channels, height, width)（如MNIST的(64,1,28,28)）。
內存佔用過高
解決：減小batch_size（如從128改爲64），或降低num_workers數量。

六、總結¶

Dataset和DataLoader是PyTorch數據處理的基石。Dataset負責“定義數據存儲規則”，DataLoader負責“高效批量讀取數據”。掌握它們的核心邏輯後，無論處理張量數據、圖像數據還是自定義數據，都能快速上手。

關鍵步驟回顧：
1. 定義Dataset（或使用TensorDataset）存儲數據。
2. 創建DataLoader，設置batch_size、shuffle等參數。
3. 迭代DataLoader，獲取批次數據並輸入模型訓練。

通過本文的實戰練習，你已能獨立完成數據加載的基本流程。後續可嘗試處理更復雜的數據（如CSV表格、自定義圖像文件夾），進一步熟悉Dataset的繼承與擴展。