PyTorch 数据加载全攻略：从自定义数据集到模型训练
目录
一、为什么需要数据加载器？
二、自定义 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. 代码解析
- 数据生成：用
make_regression生成带噪声的回归数据，并转换为 PyTorch 张量。 - 数据集构建：用
TensorDataset将特征和标签绑定，方便后续加载。 - 批量加载：
DataLoader按批次读取数据，每次训练用 100 个样本。 - 模型训练：线性回归模型通过梯度下降优化，最终输出预测的权重和偏置，与真实值对比。
六、总结与拓展
本文介绍了 PyTorch 中数据加载的核心工具：
- 自定义 Dataset：灵活处理特殊数据格式
- TensorDataset：快速包装张量数据
- DataLoader：高效批量加载，支持多线程和数据打乱
在实际项目中，你可以根据数据类型选择合适的工具：
- 处理图片：用
ImageFolder（PyTorch 内置，支持按文件夹分类） - 处理文本：自定义 Dataset 读取文本文件并转换为张量
- 大规模数据：结合
num_workers和pin_memory（针对 GPU 加速）
掌握数据加载是深度学习的基础，用好这些工具能让你的训练流程更高效、更易维护。快去试试用它们处理你的数据吧！