在 PyTorch 中，Dataset 和 DataLoader 是处理数据的核心工具。它们的作用是将数据高效地加载到模型中，支持批量处理、多线程加速和数据增强等功能。
一、Dataset：数据集的抽象
Dataset 是一个抽象类，用于表示数据集的接口。你需要继承 torch.utils.data.Dataset 并实现以下两个方法：
__len__(): 返回数据集的总样本数。__getitem__(idx): 根据索引idx返回一个样本（数据和标签）。
示例：自定义 Dataset
import torch
from torch.utils.data import Datasetclass CustomDataset(Dataset):def __init__(self, data, labels, transform=None):self.data = dataself.labels = labelsself.transform = transform # 数据预处理/增强函数def __len__(self):return len(self.data)def __getitem__(self, idx):sample = {"data": self.data[idx], "label": self.labels[idx]}if self.transform:sample = self.transform(sample)return sample
使用场景
- 加载图像、文本、表格数据等。
- 支持数据预处理（如归一化、裁剪）和数据增强（如随机翻转）。
二、 DataLoader：高效加载数据
DataLoader 负责将 Dataset 包装成一个可迭代对象，支持批量加载、多线程加速和数据打乱。
基本用法
from torch.utils.data import DataLoader# 假设 dataset 是你的 CustomDataset 实例
data_loader = DataLoader(dataset,batch_size=32, # 批量大小shuffle=True, # 是否打乱数据（训练时建议开启）num_workers=4, # 多线程加载数据的进程数drop_last=False # 是否丢弃最后不足一个 batch 的数据
)
遍历 DataLoader
for batch in data_loader:data = batch["data"] # 形状：[batch_size, ...]labels = batch["label"] # 形状：[batch_size]# 将数据送入模型训练...
三、pytorch内置数据集
PyTorch 提供了一系列内置数据集，这些数据集可以直接用于训练模型。这些数据集涵盖了多种领域，如图像、文本、音频等。以下是一些常用的PyTorch内置数据集：
图像数据集
-
MNIST: 手写数字数据集，包含0到9的手写数字图片。
from torchvision import datasets mnist_train = datasets.MNIST(root='./data', train=True, download=True, transform=transform) -
CIFAR10/CIFAR100: 包含彩色图片的数据集，CIFAR10有60000张32x32的彩色图片，分为10个类别；CIFAR100类似但有100个类别。
cifar10_train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) -
ImageNet: 包含超过1400万张图片的非常庞大的数据集，常用于图像识别和分类任务。
import torchvision.datasets as datasets imagenet_train = datasets.ImageNet(root='./data', split='train', download=True) -
STL10: 一个用于计算机视觉研究的小型图像数据集，包含96x96的彩色图片。
stl10_train = datasets.STL10(root='./data', split='train', download=True) -
SVHN: 包含数字图片的数据集，与MNIST类似但包含更多实际场景的图片。
svhn_train = datasets.SVHN(root='./data', split='train', download=True, transform=transform)
文本数据集
1.Text8: 一个用于自然语言处理的小型文本数据集。
from torchtext.datasets import Text8
text8_train = Text8(split=('train',))2. AG_NEWS: 包含新闻文章的文本数据集，分为4个类别。
from torchtext.datasets import AG_NEWS
ag_news_train = AG_NEWS(split=('train',))音频数据集
1. Speech Commands: 一个用于语音识别的数据集，包含约65,000个单词发音的音频文件。
from torchaudio.datasets import SPEECHCOMMANDS
speech_commands = SPEECHCOMMANDS(root="./data", download=True) 使用方法
要使用这些数据集，首先需要导入torchvision（对于图像数据集）、torchtext（对于文本数据集）或torchaudio（对于音频数据集），然后使用其提供的类来加载数据。通常还包括一些数据预处理步骤，例如转换（transforms）。
import torchvision.transforms as transforms
from torchvision import datasetstransform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
mnist_train = datasets.MNIST(root='./data', train=True, download=True, transform=transform)四、完整代码示例
步骤 1：创建数据集
import numpy as np
from torch.utils.data import Dataset, DataLoader# 生成示例数据（假设是 10 个样本，每个样本是长度为 5 的向量）
data = np.random.randn(10, 5)
labels = np.random.randint(0, 2, size=(10,)) # 二分类标签class MyDataset(Dataset):def __init__(self, data, labels):self.data = torch.tensor(data, dtype=torch.float32)self.labels = torch.tensor(labels, dtype=torch.long)def __len__(self):return len(self.data)def __getitem__(self, idx):return {"data": self.data[idx],"label": self.labels[idx]}dataset = MyDataset(data, labels)
步骤 2：创建 DataLoader
data_loader = DataLoader(dataset,batch_size=2,shuffle=True,num_workers=2
)
步骤 3：使用 DataLoader 训练模型
model = ... # 你的模型
optimizer = torch.optim.Adam(model.parameters())
loss_fn = torch.nn.CrossEntropyLoss()for epoch in range(10):for batch in data_loader:x = batch["data"]y = batch["label"]# 前向传播outputs = model(x)loss = loss_fn(outputs, y)# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()
五、常见问题解决
（1）数据格式不匹配
- 问题：
DataLoader返回的数据形状与模型输入不匹配。 - 解决：检查
Dataset的__getitem__返回的数据类型和形状，确保与模型输入一致。
（2）多线程加载卡顿
- 问题：设置
num_workers>0时程序卡死或报错。 - 解决：在 Windows 系统中，多线程可能需要将代码放在
if __name__ == "__main__":块中运行。
（3）数据增强
- 使用
torchvision.transforms中的工具（如RandomCrop、RandomHorizontalFlip）对图像数据进行增强：from torchvision import transformstransform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5], std=[0.5]), ])
（4）内存不足
- 对于大型数据集，使用
torch.utils.data.DataLoader的persistent_workers=True（PyTorch 1.7+）或优化数据加载逻辑。
六、高级功能
- 分布式训练：使用
torch.utils.data.distributed.DistributedSampler配合多 GPU。 - 预加载数据：使用
torch.utils.data.TensorDataset直接加载 Tensor 数据。 - 自定义采样器：通过
sampler参数控制数据采样顺序（如平衡类别采样）。