目錄

一、為什么需要數據加載器？

二、自定義 Dataset 類

1. 核心方法解析

2. 代碼實現

三、快速上手：TensorDataset

1. 代碼示例

2. 適用場景

四、DataLoader：批量加載數據的利器

1. 核心參數說明

2. 代碼示例

五、實戰：用數據加載器訓練線性回歸模型

1. 完整代碼

2. 代碼解析

六、總結與拓展

---

在深度學習實踐中，數據加載是模型訓練的第一步，也是至關重要的一環。高效的數據加載不僅能提高訓練效率，還能讓代碼更具可維護性。本文將結合 PyTorch 的核心 API，通過實例詳解數據加載的全過程，從自定義數據集到批量訓練，帶你快速掌握 PyTorch 數據處理的精髓。

一、為什么需要數據加載器？

在處理大規模數據時，我們不可能一次性將所有數據加載到內存中。PyTorch 提供了Dataset和DataLoader兩個核心類來解決這個問題：

• Dataset：負責數據的存儲和索引
• DataLoader：負責批量加載、打亂數據和多線程處理

簡單來說，Dataset就像一個 "倉庫"，而DataLoader是 "搬運工"，負責把數據按批次運送到模型中進行訓練。

二、自定義 Dataset 類

當我們需要處理特殊格式的數據（如自定義標注文件、特殊預處理）時，就需要自定義數據集。自定義數據集需繼承torch.utils.data.Dataset，并實現三個核心方法：

1. 核心方法解析

• __init__：初始化數據集，加載數據路徑或原始數據
• __len__：返回數據集的樣本數量
• __getitem__：根據索引返回單個樣本（特征 + 標簽）

2. 代碼實現

import torch
from torch.utils.data import Datasetclass MyDataset(Dataset):def __init__(self, data, labels):# 初始化數據和標簽self.data = dataself.labels = labelsdef __len__(self):# 返回樣本總數return len(self.data)def __getitem__(self, index):# 根據索引返回單個樣本sample = self.data[index]label = self.labels[index]return sample, label# 使用示例
if __name__ == "__main__":# 生成隨機數據x = torch.randn(1000, 100, dtype=torch.float32)  # 1000個樣本，每個100個特征y = torch.randn(1000, 1, dtype=torch.float32)   # 對應的標簽# 創建自定義數據集dataset = MyDataset(x, y)print(f"數據集大小：{len(dataset)}")print(f"第一個樣本：{dataset[0]}")  # 查看第一個樣本

三、快速上手：TensorDataset

如果你的數據已經是 PyTorch 張量（Tensor），且不需要復雜的預處理，那么TensorDataset會是更好的選擇。它是 PyTorch 內置的數據集類，能快速將特征和標簽綁定在一起。

1. 代碼示例

from torch.utils.data import TensorDataset, DataLoader# 生成張量數據
x = torch.randn(1000, 100, dtype=torch.float32)
y = torch.randn(1000, 1, dtype=torch.float32)# 使用TensorDataset包裝數據
dataset = TensorDataset(x, y)  # 特征和標簽按索引對應# 查看樣本
print(f"樣本數量：{len(dataset)}")
print(f"第一個樣本特征：{dataset[0][0].shape}")
print(f"第一個樣本標簽：{dataset[0][1]}")

2. 適用場景

• 數據已轉換為 Tensor 格式
• 不需要復雜的預處理邏輯
• 快速搭建訓練流程（如驗證代碼可行性）

四、DataLoader：批量加載數據的利器

有了數據集，還需要高效的批量加載工具。DataLoader可以實現：

• 批量讀取數據（batch_size）
• 打亂數據順序（shuffle）
• 多線程加載（num_workers）

1. 核心參數說明

參數	作用
dataset	要加載的數據集
batch_size	每批樣本數量（常用 32/64/128）
shuffle	每個 epoch 是否打亂數據（訓練時設為 True）
num_workers	加載數據的線程數（加速數據讀取）

2. 代碼示例

# 創建DataLoader
dataloader = DataLoader(dataset=dataset,batch_size=32,      # 每批32個樣本shuffle=True,       # 訓練時打亂數據num_workers=2       # 2個線程加載
)# 遍歷數據
for batch_idx, (batch_x, batch_y) in enumerate(dataloader):print(f"第{batch_idx}批：")print(f"特征形狀：{batch_x.shape}")  # (32, 100)print(f"標簽形狀：{batch_y.shape}")  # (32, 1)if batch_idx == 2:  # 只看前3批break

五、實戰：用數據加載器訓練線性回歸模型

下面結合一個完整案例，展示如何使用TensorDataset和DataLoader訓練模型。我們將實現一個線性回歸任務，預測生成的隨機數據。

1. 完整代碼

from sklearn.datasets import make_regression
import torch
from torch.utils.data import TensorDataset, DataLoader
from torch import nn, optim# 生成回歸數據
def build_data():bias = 14.5# 生成1000個樣本，100個特征x, y, coef = make_regression(n_samples=1000,n_features=100,n_targets=1,bias=bias,coef=True,random_state=0  # 固定隨機種子，保證結果可復現)# 轉換為Tensor并調整形狀x = torch.tensor(x, dtype=torch.float32)y = torch.tensor(y, dtype=torch.float32).view(-1, 1)  # 轉為列向量bias = torch.tensor(bias, dtype=torch.float32)coef = torch.tensor(coef, dtype=torch.float32)return x, y, coef, bias# 訓練函數
def train():x, y, true_coef, true_bias = build_data()# 構建數據集和數據加載器dataset = TensorDataset(x, y)dataloader = DataLoader(dataset=dataset,batch_size=100,  # 每批100個樣本shuffle=True     # 訓練時打亂數據)# 定義模型、損失函數和優化器model = nn.Linear(in_features=x.size(1), out_features=y.size(1))  # 線性層criterion = nn.MSELoss()  # 均方誤差損失optimizer = optim.SGD(model.parameters(), lr=0.01)  # 隨機梯度下降# 訓練50個epochepochs = 50for epoch in range(epochs):for batch_x, batch_y in dataloader:# 前向傳播y_pred = model(batch_x)loss = criterion(batch_y, y_pred)# 反向傳播和參數更新optimizer.zero_grad()  # 清空梯度loss.backward()        # 計算梯度optimizer.step()       # 更新參數# 打印結果print(f"真實權重：{true_coef[:5]}...")  # 只顯示前5個print(f"預測權重：{model.weight.detach().numpy()[0][:5]}...")print(f"真實偏置：{true_bias}")print(f"預測偏置：{model.bias.item()}")if __name__ == "__main__":train()

2. 代碼解析

1. 數據生成：用make_regression生成帶噪聲的回歸數據，并轉換為 PyTorch 張量。
2. 數據集構建：用TensorDataset將特征和標簽綁定，方便后續加載。
3. 批量加載：DataLoader按批次讀取數據，每次訓練用 100 個樣本。
4. 模型訓練：線性回歸模型通過梯度下降優化，最終輸出預測的權重和偏置，與真實值對比。

六、總結與拓展

本文介紹了 PyTorch 中數據加載的核心工具：

• 自定義 Dataset：靈活處理特殊數據格式
• TensorDataset：快速包裝張量數據
• DataLoader：高效批量加載，支持多線程和數據打亂

在實際項目中，你可以根據數據類型選擇合適的工具：

• 處理圖片：用ImageFolder（PyTorch 內置，支持按文件夾分類）
• 處理文本：自定義 Dataset 讀取文本文件并轉換為張量
• 大規模數據：結合num_workers和pin_memory（針對 GPU 加速）

掌握數據加載是深度學習的基礎，用好這些工具能讓你的訓練流程更高效、更易維護。快去試試用它們處理你的數據吧！