PyTorch Dataset与DataLoader实战：自定义数据集与参数配置详解
在PyTorch中，Dataset和DataLoader是构建数据管道的核心组件。Dataset封装数据源及访问方式，DataLoader负责批量加载、打乱、多进程并行等。本文将从自定义数据集实现、内置数据集加载、DataLoader参数详解到实战案例，帮助读者搭建高效的数据流水线。## 自定义数据集：映射式与可迭代式
自定义Dataset必须继承torch.utils.data.Dataset并实现三个方法：
- __init__：初始化数据路径、标签、变换等；
- __len__：返回数据集样本总数；
- __getitem__(self, idx)：根据索引返回样本（特征和标签）。
示例：
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, idx):
return self.data, self.labels
PyTorch支持两种数据集风格：
- 映射式（Map-Style）：实现__getitem__和__len__，通过索引随机访问，适用于数据能全部放入内存索引结构的场景，支持shuffle。
- 可迭代式（Iterable-Style）：继承IterableDataset，实现__iter__返回迭代器，适用于数据流式读取或无法统计长度的场景，不能使用len()和shuffle。
## 内置数据集加载
torchvision.datasets提供多种内置数据集，如FashionMNIST。加载时需指定root路径、train标志、transform（如ToTensor）和download是否自动下载。
import torch
from torchvision import datasets
from torchvision.transforms import ToTensor
training_data = datasets.FashionMNIST(
root="./data",
train=True,
download=True,
transform=ToTensor()
)
test_data = datasets.FashionMNIST(
root="./data",
train=False,
download=True,
transform=ToTensor()
)
访问样本：training_data返回(图像张量, 标签)，标签映射为0~9对应的服装类别。
## DataLoader参数详解
DataLoader将DataSet封装为迭代器，核心API及参数含义：
DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
batch_sampler=None, num_workers=0, collate_fn=None,
pin_memory=False, drop_last=False, timeout=0,
worker_init_fn=None, generator=None,
prefetch_factor=2, persistent_workers=False)
- dataset：需要加载的Dataset对象。
- batch_size：每个批次样本数，默认1。
- shuffle：每个epoch开始时是否打乱数据（仅映射式有效，使用RandomSampler）。
- sampler：自定义索引采样器，若指定则shuffle必须为False。
- batch_sampler：返回批次索引，与batch_size、shuffle、sampler互斥。
- num_workers：数据加载子进程数，0表示主进程，>0可加速。
- collate_fn：将样本列表合并成批次，默认按第0维堆叠；样本结构不一致时需自定义。
- pin_memory：将张量复制到CUDA固定内存，加速CPU到GPU传输（仅CUDA）。
- drop_last：丢弃最后一个不完整批次，BatchNorm等需固定批次大小时建议True。
- timeout：获取批次超时时间（秒）。
- worker_init_fn：worker初始化函数，常用于设置随机种子。
- generator：伪随机数生成器，保证可复现。
- prefetch_factor：每个worker预加载的batch数（默认2）。
- persistent_workers：是否保持worker进程存活以加速后续epoch。
## 实战案例：鸢尾花数据集
使用sklearn的iris数据集，实现数据归一化、随机打乱，并划分训练/验证/测试集，通过DataLoader批量加载。
import torch
from sklearn.datasets import load_iris
from torch.utils.data import Dataset, DataLoader
def load_data(shuffle=True):
x = torch.tensor(load_iris().data)
y = torch.tensor(load_iris().target)
# 数据归一化
x_min = torch.min(x, dim=0).values
x_max = torch.max(x, dim=0).values
x = (x - x_min) / (x_max - x_min)
if shuffle:
idx = torch.randperm(x.shape)
x = x
y = y
return x, y
class IrisDataset(Dataset):
def __init__(self, mode='train', num_train=120, num_dev=15):
super().__init__()
x, y = load_data(shuffle=True)
if mode == 'train':
self.x, self.y = x[:num_train], y[:num_train]
elif mode == 'dev':
self.x, self.y = x, y
else:
self.x, self.y = x, y
def __getitem__(self, idx):
return self.x, self.y
def __len__(self):
return len(self.x)
batch_size = 16
train_dataset = IrisDataset(mode='train')
dev_dataset = IrisDataset(mode='dev')
test_dataset = IrisDataset(mode='test')
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
dev_loader = DataLoader(dev_dataset, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True)
# 遍历训练集
for batch_x, batch_y in train_loader:
# 训练模型
pass
## 总结
- Dataset定义数据源与访问方式，映射式最常用，流式数据用IterableDataset。
- DataLoader封装采样、批处理、多进程加载、内存固定等，参数丰富，理解每个参数可应对不同场景。
- 通过自定义sampler和collate_fn可灵活处理非平衡数据或异构样本。
- 多进程加载（num_workers>0）是加速训练的关键，需关注内存复制和系统兼容性。
掌握Dataset和DataLoader的内部机制，能构建高效数据管道，降低I/O瓶颈，充分释放GPU算力。
Re: PyTorch Dataset与DataLoader实战：自定义数据集与参数配置详解
干货满满的帖！自定义数据集的映射式和可迭代式区别讲得很清楚，DataLoader参数列表也很有参考价值。 不过实战案例里有个小笔误：`load_data` 函数里生成 `idx` 后，`x = x; y = y` 并没有真正打乱，应该改成 `x = x; y = y`？另外最后一行 `if mode == 'trai` 似乎被截断了，方便的话补全一下？期待后续完善的内容～Re: PyTorch Dataset与DataLoader实战：自定义数据集与参数配置详解
您好，感谢分享！非常详细的教程，把Dataset和DataLoader的核心概念和用法都讲透了。我尤其喜欢参数详解部分，prefetch_factor和persistent_workers这些细节平时容易忽略。提醒一下，实战案例中的load_data函数里shuffle操作好像没有正确更新x和y，应该是 x = x; y = y，目前代码中x = x, y = y没有起到打乱作用。另外，能否补充一个可迭代式数据集的具体应用场景例子？期待后续更多分享！Re: PyTorch Dataset与DataLoader实战：自定义数据集与参数配置详解
感谢楼主的分享！讲解非常细致，尤其是DataLoader各项参数的说明，平时容易忽略的`prefetch_factor`和`persistent_workers`都解释清楚了，很实用。 在看鸢尾花数据集案例时发现一个小笔误：`load_data`函数中随机打乱后，`x = x`和`y = y`似乎没实际应用`idx`，应该是`x = x`和`y = y`？另外`IrisDataset`的mode参数写成了`'trai`，可能少了个n。整体不影响理解，但顺手提一下方便其他读者。 总之，这篇实战对新手搭建数据管道很有价值，收藏了！
页:
[1]