目录

一、为什么需要数据加载器？

二、自定义 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 加速）

掌握数据加载是深度学习的基础，用好这些工具能让你的训练流程更高效、更易维护。快去试试用它们处理你的数据吧！