数据加载#
数据是深度学习的核心。PyTorch 提供了强大的数据加载和预处理工具，帮助我们高效地处理各种类型的数据。
Dataset 和 DataLoader#
在 PyTorch 中，处理和加载数据是深度学习训练过程中的关键步骤。PyTorch 提供了强大的工具，包括 torch.utils.data.Dataset 和 torch.utils.data.DataLoader，帮助我们管理数据集、批量加载和数据增强等任务。
Dataset 基类#
Dataset 是一个抽象类，所有自定义数据集都应该继承它。我们需要实现以下两个方法：
__getitem__(idx)：通过索引返回一个样本__len__()：返回数据集中的样本数量
from torch.utils.data import Dataset
class CustomDataset(Dataset):
def __init__(self, data, labels):
"""
初始化数据集
data: 输入特征
labels: 目标标签
"""
self.data = data
self.labels = labels
def __len__(self):
"""返回数据集的大小"""
return len(self.data)
def __getitem__(self, idx):
"""返回指定索引的数据"""
sample = self.data[idx]
label = self.labels[idx]
return sample, label
# 示例数据
X_data = [[1, 2], [3, 4], [5, 6], [7, 8]] # 输入特征
Y_data = [1, 0, 1, 0] # 目标标签
# 创建数据集实例
dataset = CustomDataset(X_data, Y_data)
print(f"数据集大小: {len(dataset)}")完整示例#
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
class MyDataset(Dataset):
def __init__(self, data, labels, transform=None):
self.data = data
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
sample = self.data[idx]
label = self.labels[idx]
# 转换为张量
sample = torch.FloatTensor(sample)
label = torch.LongTensor([label])[0]
# 应用变换
if self.transform:
sample = self.transform(sample)
return sample, label
# 创建数据集
data = np.random.randn(1000, 784)
labels = np.random.randint(0, 10, 1000)
dataset = MyDataset(data, labels)
# 创建 DataLoader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 使用
for batch_data, batch_labels in dataloader:
print(batch_data.shape) # (32, 784)
print(batch_labels.shape) # (32,)
breakDataLoader 参数#
DataLoader 是 PyTorch 提供的一个重要工具，用于从 Dataset 中按批次（batch）加载数据。DataLoader 允许我们批量读取数据并进行多线程加载，从而提高训练效率。
from torch.utils.data import DataLoader
dataloader = DataLoader(
dataset,
batch_size=32, # 批次大小，每次加载的样本数量
shuffle=True, # 是否打乱数据，通常训练时需要将数据打乱
num_workers=4, # 数据加载的进程数
pin_memory=True, # 是否将数据固定在内存中（GPU加速）
drop_last=False, # 如果数据集中的样本数不能被batch_size整除，是否丢弃最后一个不完整的批次
collate_fn=None # 自定义批处理函数
)
# 使用 DataLoader 迭代数据
for batch_idx, (inputs, labels) in enumerate(dataloader):
print(f'Batch {batch_idx + 1}:')
print(f'Inputs shape: {inputs.shape}')
print(f'Labels shape: {labels.shape}')
break # 只打印第一个批次重要参数说明：
batch_size：每次加载的样本数量shuffle：是否对数据进行洗牌，通常训练时需要将数据打乱num_workers：数据加载的进程数，设置为 0 表示在主进程中加载pin_memory：是否将数据固定在内存中，可以加速 GPU 传输drop_last：如果数据集中的样本数不能被batch_size整除，设置为True时丢弃最后一个不完整的 batch
图像数据集#
对于图像数据集，torchvision.datasets 提供了许多常见数据集（如 CIFAR-10、ImageNet、MNIST 等）以及用于加载图像数据的工具。
使用 torchvision 数据集#
import torchvision
import torchvision.datasets as datasets
from torchvision import transforms
# 定义变换
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# 加载 MNIST 数据集
train_dataset = datasets.MNIST(
root='./data',
train=True,
download=True, # 如果数据集不存在，自动下载
transform=transform
)
test_dataset = datasets.MNIST(
root='./data',
train=False,
download=True,
transform=transform
)
# 创建 DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# 迭代训练数据
for inputs, labels in train_loader:
print(f'Inputs shape: {inputs.shape}') # 每个批次的输入数据形状
print(f'Labels shape: {labels.shape}') # 每个批次的标签形状
break加载 CIFAR-10 数据集#
# 定义预处理操作
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
# 下载并加载 CIFAR-10 数据集
train_dataset = datasets.CIFAR10(
root='./data',
train=True,
download=True,
transform=transform
)
test_dataset = datasets.CIFAR10(
root='./data',
train=False,
download=True,
transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)合并多个数据集#
如果你的数据集由多个文件、多个来源（例如多个图像文件夹）组成，可以使用 ConcatDataset 来连接多个数据集：
from torch.utils.data import ConcatDataset
# 假设 dataset1 和 dataset2 是两个 Dataset 对象
combined_dataset = ConcatDataset([dataset1, dataset2])
combined_loader = DataLoader(combined_dataset, batch_size=64, shuffle=True)自定义图像数据集#
from PIL import Image
import os
class ImageDataset(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.images = []
self.labels = []
# 假设目录结构为：root_dir/class_name/image.jpg
for class_name in os.listdir(root_dir):
class_dir = os.path.join(root_dir, class_name)
if os.path.isdir(class_dir):
for img_name in os.listdir(class_dir):
img_path = os.path.join(class_dir, img_name)
self.images.append(img_path)
self.labels.append(class_name)
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
img_path = self.images[idx]
label = self.labels[idx]
# 加载图像
image = Image.open(img_path).convert('RGB')
# 应用变换
if self.transform:
image = self.transform(image)
return image, label数据变换（Transforms）#
数据预处理和增强对于提高模型的性能至关重要。PyTorch 提供了 torchvision.transforms 模块来进行常见的图像预处理和增强操作，如旋转、裁剪、归一化等。
常用图像变换#
from torchvision import transforms
from PIL import Image
# 组合多个变换
transform = transforms.Compose([
# 调整大小
transforms.Resize((256, 256)), # 将图像调整为指定大小
transforms.RandomResizedCrop(224), # 随机裁剪并调整为指定大小
# 数据增强
transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转
transforms.RandomVerticalFlip(p=0.5), # 随机垂直翻转
transforms.RandomRotation(degrees=15), # 随机旋转
transforms.ColorJitter(brightness=0.2, contrast=0.2), # 随机调整亮度和对比度
# 转换为张量
transforms.ToTensor(), # 将图像转换为 PyTorch 张量，值会被归一化到 [0, 1] 范围
# 归一化（标准化）
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]) # 标准化图像数据，通常使用预训练模型时需要进行标准化处理
])
# 单独使用
resize = transforms.Resize((224, 224))
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize(mean=[0.5], std=[0.5])
# 使用示例
image = Image.open('image.jpg')
image_tensor = transform(image)
print(f'Image tensor shape: {image_tensor.shape}')常用变换说明：
Resize：调整图像大小ToTensor：将图像转换为 PyTorch 张量，值会被归一化到 [0, 1] 范围Normalize：标准化图像数据，通常使用预训练模型时需要进行标准化处理Compose：将多个变换操作组合在一起
自定义变换#
class RandomNoise(object):
def __init__(self, noise_factor=0.1):
self.noise_factor = noise_factor
def __call__(self, tensor):
noise = torch.randn_like(tensor) * self.noise_factor
return tensor + noise
transform = transforms.Compose([
transforms.ToTensor(),
RandomNoise(noise_factor=0.1)
])文本数据集#
文本分类数据集#
class TextDataset(Dataset):
def __init__(self, texts, labels, tokenizer, max_length=128):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = self.labels[idx]
# 分词和编码
encoding = self.tokenizer(
text,
truncation=True,
padding='max_length',
max_length=self.max_length,
return_tensors='pt'
)
return {
'input_ids': encoding['input_ids'].flatten(),
'attention_mask': encoding['attention_mask'].flatten(),
'label': torch.tensor(label, dtype=torch.long)
}数据采样器（Sampler）#
随机采样#
from torch.utils.data import RandomSampler, SequentialSampler
# 随机采样
random_sampler = RandomSampler(dataset, replacement=False, num_samples=None)
# 顺序采样
sequential_sampler = SequentialSampler(dataset)
dataloader = DataLoader(dataset, sampler=random_sampler, batch_size=32)加权采样#
from torch.utils.data import WeightedRandomSampler
# 为每个样本分配权重
weights = [0.1, 0.9, 0.5, ...] # 对应每个样本的权重
weighted_sampler = WeightedRandomSampler(
weights=weights,
num_samples=len(weights),
replacement=True
)
dataloader = DataLoader(dataset, sampler=weighted_sampler, batch_size=32)批次采样#
from torch.utils.data import BatchSampler
batch_sampler = BatchSampler(
sampler=RandomSampler(dataset),
batch_size=32,
drop_last=False
)
dataloader = DataLoader(dataset, batch_sampler=batch_sampler)多进程数据加载#
# 设置 num_workers > 0 启用多进程
dataloader = DataLoader(
dataset,
batch_size=32,
shuffle=True,
num_workers=4, # 使用4个进程加载数据
pin_memory=True, # 加速 GPU 传输
persistent_workers=True # 保持 worker 进程存活
)注意：
- Windows 上可能需要将代码放在
if __name__ == '__main__':中 num_workers=0表示在主进程中加载数据- 太多 worker 可能导致内存问题
数据预处理技巧#
1. 数据归一化#
# 计算数据集的均值和标准差
def compute_mean_std(dataset):
mean = torch.zeros(3)
std = torch.zeros(3)
total_samples = 0
for data, _ in dataset:
batch_samples = data.size(0)
data = data.view(batch_samples, data.size(1), -1)
mean += data.mean(2).sum(0)
std += data.std(2).sum(0)
total_samples += batch_samples
mean /= total_samples
std /= total_samples
return mean, std2. 数据增强策略#
# 训练时的增强
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# 验证/测试时的变换（不增强）
val_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])3. 处理不平衡数据#
from torch.utils.data import WeightedRandomSampler
# 计算每个类别的权重
class_counts = [100, 50, 200, 30] # 每个类别的样本数
total_samples = sum(class_counts)
class_weights = [total_samples / count for count in class_counts]
# 为每个样本分配权重
sample_weights = [class_weights[label] for label in labels]
sampler = WeightedRandomSampler(
weights=sample_weights,
num_samples=len(sample_weights),
replacement=True
)性能优化#
1. 预加载数据#
# 在 __init__ 中预加载所有数据到内存
class PreloadedDataset(Dataset):
def __init__(self, data_paths):
self.data = []
for path in data_paths:
self.data.append(self.load_data(path))
def __getitem__(self, idx):
return self.data[idx]2. 使用 pin_memory#
dataloader = DataLoader(
dataset,
batch_size=32,
pin_memory=True, # 将数据固定在内存中，加速 GPU 传输
num_workers=4
)3. 合理设置 batch_size#
# 根据 GPU 内存调整 batch_size
# 较大的 batch_size 可以提高训练速度，但需要更多内存
batch_size = 32 # 根据实际情况调整常见问题#
1. 内存不足#
# 解决方案：减小 batch_size 或使用更少的 workers
dataloader = DataLoader(dataset, batch_size=16, num_workers=2)2. 数据加载慢#
# 解决方案：增加 workers 或使用 pin_memory
dataloader = DataLoader(
dataset,
batch_size=32,
num_workers=8,
pin_memory=True
)3. Windows 多进程问题#
# 在 Windows 上，将主代码放在 if __name__ == '__main__' 中
if __name__ == '__main__':
dataloader = DataLoader(dataset, num_workers=4)练习#
- 创建一个自定义的图像数据集类
- 实现一个文本分类的数据集类
- 使用数据增强技术提高模型泛化能力
- 处理不平衡数据集