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的继承与扩展。