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 Dataset

class MyDataset(Dataset):
    def __init__(self, data, labels):
        # 初始化数据和标签
        self.data = data
        self.labels = labels
        
    def __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. 核心参数说明

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个epoch
    epochs = 50
    for 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 加速）

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