当前位置：首页 > news >正文

吴恩达机器学习课程（PyTorch 适配）学习笔记：2.3 PyTorch 工具与高效实现

news 2025/10/9 10:05:30

PyTorch 提供了丰富的工具和API，使得神经网络的实现既简洁又高效。掌握这些工具的使用方法和高效实现的技巧，对于构建高性能深度学习模型至关重要。本章将介绍神经网络的高效实现思路、矩阵乘法的最佳实践以及PyTorch的数据处理流水线。

2.3.1 神经网络高效实现思路

高效实现神经网络不仅能加快模型训练速度，还能提高代码的可读性和可维护性。以下是一些关键的高效实现思路：

模块化设计原则

模块化是高效实现神经网络的基础，遵循以下原则：

单一职责：每个模块只负责一个特定功能
可复用性：设计通用模块，可在不同网络中复用
可组合性：模块之间可以灵活组合形成复杂网络
可扩展性：便于添加新功能或修改现有功能

利用PyTorch内置组件

PyTorch提供了许多经过优化的内置组件，应优先使用：

nn模块：包含各种预定义的层和激活函数
nn.functional：提供各种功能函数，如激活函数、损失函数
nn.init：提供多种参数初始化方法
optim：提供各种优化器实现

计算图优化

PyTorch的动态计算图虽然灵活，但也需要注意优化：

减少不必要的操作：避免冗余计算和内存分配
使用in-place操作：适当使用带_后缀的in-place操作（如relu_()）减少内存占用
合理使用中间变量：避免创建不必要的中间张量
利用广播机制：减少显式的维度扩展操作

并行计算策略

充分利用硬件资源进行并行计算：

数据并行（Data Parallelism）：在多个GPU上分割数据
模型并行（Model Parallelism）：将模型不同部分分配到不同GPU
自动混合精度（AMP）：使用半精度浮点数加速计算
多线程数据加载：利用DataLoader的多进程数据加载

高效实现代码示例

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset# 1. 模块化网络组件设计
class ConvBlock(nn.Module):"""卷积块：卷积 + 批归一化 + 激活函数"""def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):super(ConvBlock, self).__init__()self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)self.bn = nn.BatchNorm2d(out_channels)self.relu = nn.ReLU(inplace=True)  # 使用in-place操作节省内存def forward(self, x):x = self.conv(x)x = self.bn(x)x = self.relu(x)return xclass ResidualBlock(nn.Module):"""残差块：两个卷积块 + 跳跃连接"""def __init__(self, channels):super(ResidualBlock, self).__init__()self.conv1 = ConvBlock(channels, channels)self.conv2 = ConvBlock(channels, channels)def forward(self, x):residual = xx = self.conv1(x)x = self.conv2(x)x += residual  # 跳跃连接return x# 2. 高效网络组装
class EfficientNet(nn.Module):def __init__(self, num_classes=10):super(EfficientNet, self).__init__()# 特征提取部分self.features = nn.Sequential(ConvBlock(3, 64, kernel_size=7, stride=2, padding=3),nn.MaxPool2d(2, 2),ResidualBlock(64),ResidualBlock(64),ConvBlock(64, 128, stride=2),ResidualBlock(128),ResidualBlock(128),nn.AdaptiveAvgPool2d((1, 1))  # 自适应池化，兼容不同输入尺寸)# 分类部分self.classifier = nn.Sequential(nn.Flatten(),nn.Linear(128, num_classes))# 初始化参数self._initialize_weights()def forward(self, x):x = self.features(x)x = self.classifier(x)return xdef _initialize_weights(self):"""自定义参数初始化"""for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')if m.bias is not None:nn.init.constant_(m.bias, 0)elif isinstance(m, nn.BatchNorm2d):nn.init.constant_(m.weight, 1)nn.init.constant_(m.bias, 0)elif isinstance(m, nn.Linear):nn.init.normal_(m.weight, 0, 0.01)nn.init.constant_(m.bias, 0)# 3. 高效训练循环
def train_model(model, train_loader, criterion, optimizer, device, num_epochs=10):model.train()model.to(device)# 使用自动混合精度scaler = torch.cuda.amp.GradScaler() if device.type == 'cuda' else Nonefor epoch in range(num_epochs):running_loss = 0.0for inputs, labels in train_loader:inputs, labels = inputs.to(device), labels.to(device)# 清零梯度optimizer.zero_grad()# 前向传播with torch.cuda.amp.autocast(enabled=scaler is not None):outputs = model(inputs)loss = criterion(outputs, labels)# 反向传播和优化if scaler:scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()else:loss.backward()optimizer.step()running_loss += loss.item() * inputs.size(0)epoch_loss = running_loss / len(train_loader.dataset)print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}')return model# 4. 主函数：整合所有组件
def main():# 设备配置device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')print(f'使用设备: {device}')# 模拟数据集class DummyDataset(Dataset):def __len__(self):return 1000def __getitem__(self, idx):return torch.randn(3, 224, 224), torch.randint(0, 10, (1,)).item()# 数据加载器（多进程）train_dataset = DummyDataset()train_loader = DataLoader(train_dataset,batch_size=32,shuffle=True,num_workers=4,  # 多进程加载pin_memory=True if device.type == 'cuda' else False  # 锁页内存)# 初始化模型、损失函数和优化器model = EfficientNet(num_classes=10)# 数据并行（如果有多个GPU）if device.type == 'cuda' and torch.cuda.device_count() > 1:model = nn.DataParallel(model)criterion = nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 训练模型model = train_model(model, train_loader, criterion, optimizer, device, num_epochs=5)if __name__ == '__main__':main()

注意事项与最佳实践

设备管理
- 始终显式指定设备（CPU/GPU）
- 确保模型和数据在同一设备上
- 使用torch.device统一管理设备
内存优化
- 适当使用in-place操作（但要注意自动求导的兼容性）
- 训练时清理不再需要的中间变量
- 对于大型模型，考虑使用模型并行
代码可读性
- 使用有意义的变量和函数命名
- 为复杂模块添加文档字符串
- 保持一致的代码风格
可复现性
- 设置随机种子（torch.manual_seed()）
- 记录实验参数和结果
- 确保代码在不同环境下的兼容性
性能监控
- 使用PyTorch Profiler分析性能瓶颈
- 监控GPU利用率和内存使用情况
- 定期评估训练和推理速度

2.3.2 矩阵乘法（规则 + PyTorch 代码框架）

矩阵乘法是神经网络中最核心的运算之一，理解其规则并掌握高效实现方法对提升模型性能至关重要。

矩阵乘法的数学规则

矩阵乘法是一种特殊的数学运算，遵循以下规则：

维度兼容性：
- 若矩阵A的形状为(m, n)，矩阵B的形状为(n, p)，则它们可以相乘
- 结果矩阵C的形状为(m, p)
元素计算：
- 结果矩阵C的元素C[i][j]等于A的第i行与B的第j列的点积
- 公式： $\sum_{k=0}^{n-1} A[i][k] \times B[k][j]$
运算性质：
- 不满足交换律： $\times B \neq B \times A$
- 满足结合律： $\times B) \times C = A \times (B \times C)$
- 满足分配律： $\times (B + C) = A \times B + A \times C$

矩阵乘法在神经网络中的应用

神经网络中的许多运算本质上都是矩阵乘法：

全连接层： $\times W^T + b$
- X: 输入，形状为(batch_size, input_features)
- W: 权重，形状为(output_features, input_features)
- Z: 输出，形状为(batch_size, output_features)
卷积层：可以视为多通道的滑动窗口矩阵乘法
注意力机制：查询、键、值之间的相似度计算

PyTorch中的矩阵乘法实现

PyTorch提供了多种矩阵乘法函数，适用于不同场景：

import torch
import time# 1. 基本矩阵乘法操作
def matrix_multiplication_demo():# 创建随机矩阵A = torch.randn(3, 4)  # 3x4矩阵B = torch.randn(4, 5)  # 4x5矩阵（与A兼容）print("矩阵A形状:", A.shape)print("矩阵B形状:", B.shape)# 方法1: torch.matmul (推荐)C1 = torch.matmul(A, B)print("matmul结果形状:", C1.shape)  # 应为(3, 5)# 方法2: @ 运算符 (Python 3.5+)C2 = A @ Bprint("运算符@结果形状:", C2.shape)  # 应为(3, 5)# 方法3: torch.mm (仅适用于2D矩阵)C3 = torch.mm(A, B)print("mm结果形状:", C3.shape)  # 应为(3, 5)# 验证结果一致性print("结果是否一致:", torch.allclose(C1, C2) and torch.allclose(C1, C3))# 批量矩阵乘法 (3D及以上)batch_size = 10A_batch = torch.randn(batch_size, 3, 4)  # 10个3x4矩阵B_batch = torch.randn(batch_size, 4, 5)  # 10个4x5矩阵# 批量矩阵乘法C_batch = torch.bmm(A_batch, B_batch)  # 每个批次单独相乘print("批量矩阵乘法结果形状:", C_batch.shape)  # 应为(10, 3, 5)# 广播机制下的矩阵乘法A_broad = torch.randn(2, 3, 4)B_broad = torch.randn(4, 5)C_broad = torch.matmul(A_broad, B_broad)  # B会广播到(2, 4, 5)print("广播矩阵乘法结果形状:", C_broad.shape)  # 应为(2, 3, 5)# 2. 神经网络中的矩阵乘法
def nn_matrix_multiplication():# 模拟全连接层计算batch_size = 32input_features = 128output_features = 64# 输入和权重X = torch.randn(batch_size, input_features)  # 输入W = torch.randn(output_features, input_features)  # 权重b = torch.randn(output_features)  # 偏置# 手动计算全连接层Z_manual = X @ W.T + b  # 等价于 torch.matmul(X, W.T) + bprint("手动计算全连接层输出形状:", Z_manual.shape)  # 应为(32, 64)# 使用nn.Linear验证linear = nn.Linear(input_features, output_features)Z_nn = linear(X)print("nn.Linear输出形状:", Z_nn.shape)  # 应为(32, 64)# 展示矩阵乘法性能def test_performance():sizes = [(100, 100), (500, 500), (1000, 1000), (2000, 2000)]for m, n in sizes:# 创建两个m×n和n×m的矩阵A = torch.randn(m, n).cuda() if torch.cuda.is_available() else torch.randn(m, n)B = torch.randn(n, m).cuda() if torch.cuda.is_available() else torch.randn(n, m)# 计时start = time.time()for _ in range(100):C = A @ Bend = time.time()# 计算每秒运算次数operations = 2 * m * n * m  # 每个矩阵乘法的大致运算次数total_ops = operations * 100gflops = total_ops / (1e9 * (end - start))device = "GPU" if torch.cuda.is_available() else "CPU"print(f"{device}上 {m}x{n} × {n}x{m} 矩阵乘法: {end-start:.4f}秒, {gflops:.2f} GFLOPS")test_performance()# 3. 矩阵乘法优化技巧
def matrix_multiplication_optimization():# 1. 利用矩阵转置减少计算量A = torch.randn(1000, 2000)B = torch.randn(1000, 2000)# 计算A^T × B 与 A × B^T 的效率对比start = time.time()C1 = A.T @ B  # 2000x1000 × 1000x2000 = 2000x2000t1 = time.time() - startstart = time.time()C2 = A @ B.T  # 1000x2000 × 2000x1000 = 1000x1000t2 = time.time() - startprint(f"A^T × B 耗时: {t1:.4f}秒")print(f"A × B^T 耗时: {t2:.4f}秒")print(f"速度比: {t1/t2:.2f}x")# 2. 利用连续内存提升性能A = torch.randn(1000, 1000)A_non_contiguous = A[:, ::2]  # 创建非连续张量print("A是否连续:", A.is_contiguous())print("A_non_contiguous是否连续:", A_non_contiguous.is_contiguous())B = torch.randn(500, 1000)start = time.time()for _ in range(100):C = A_non_contiguous @ B.Tt_non_contiguous = time.time() - start# 转为连续张量A_contiguous = A_non_contiguous.contiguous()start = time.time()for _ in range(100):C = A_contiguous @ B.Tt_contiguous = time.time() - startprint(f"非连续矩阵乘法耗时: {t_non_contiguous:.4f}秒")print(f"连续矩阵乘法耗时: {t_contiguous:.4f}秒")print(f"性能提升: {t_non_contiguous/t_contiguous:.2f}x")if __name__ == "__main__":print("=== 基本矩阵乘法操作 ===")matrix_multiplication_demo()print("\n=== 神经网络中的矩阵乘法 ===")nn_matrix_multiplication()if torch.cuda.is_available():print("\n=== 矩阵乘法优化技巧 ===")matrix_multiplication_optimization()

矩阵乘法的高效实现技巧

选择合适的函数：
- 2D矩阵：torch.matmul 或 @ 运算符
- 批量2D矩阵：torch.bmm
- 高维张量：torch.matmul（自动处理广播）
内存布局优化：
- 确保张量是连续的（使用.contiguous()）
- 适当调整维度顺序，使内存访问更高效
- 利用torch.transpose而非手动切片获取转置
精度选择：
- 训练时可使用float32保证精度
- 推理时可考虑float16或bfloat16加速计算
- 对精度要求不高的场景，可尝试int8量化
利用硬件特性：
- GPU上使用torch.backends.cudnn.benchmark = True
- 对于固定形状的输入，预热后再进行正式计算
- 合理设置矩阵分块大小，充分利用缓存

常见错误与避坑指南

维度不匹配：
- 错误：尝试相乘形状不兼容的矩阵
- 解决：使用.shape检查矩阵维度，确保前一个矩阵的列数等于后一个矩阵的行数
混淆矩阵乘法与元素乘法：
- 错误：使用*进行矩阵乘法（*是元素乘法）
- 解决：矩阵乘法使用@或torch.matmul，元素乘法使用*或torch.mul
忽略批量维度：
- 错误：对包含批量维度的张量使用torch.mm
- 解决：批量矩阵乘法使用torch.bmm或torch.matmul
内存溢出：
- 错误：尝试计算过大的矩阵乘法导致OOM
- 解决：减小批量大小，使用模型并行，或使用更低精度

2.3.3 PyTorch 数据处理（Tensor+Dataset+DataLoader）

高效的数据处理是深度学习 pipeline 的关键环节，PyTorch 提供了 Tensor、Dataset 和 DataLoader 等工具，使数据处理变得简洁高效。

Tensor 数据结构

Tensor 是 PyTorch 中最基本的数据结构，类似于多维数组，但支持 GPU 加速和自动求导。

Tensor 的基本操作

import torch
import numpy as np# 1. Tensor 创建
def create_tensors():# 从列表创建tensor_from_list = torch.tensor([[1, 2], [3, 4]])print("从列表创建的Tensor:\n", tensor_from_list)# 从NumPy数组创建np_array = np.array([[1, 2], [3, 4]])tensor_from_np = torch.from_numpy(np_array)print("\n从NumPy创建的Tensor:\n", tensor_from_np)# 创建特殊张量zeros_tensor = torch.zeros((2, 3))  # 全零张量ones_tensor = torch.ones((2, 3))   # 全一张量rand_tensor = torch.randn((2, 3))  # 正态分布随机张量eye_tensor = torch.eye(3)          # 单位矩阵print("\n全零张量:\n", zeros_tensor)print("全一张量:\n", ones_tensor)print("随机张量:\n", rand_tensor)print("单位矩阵:\n", eye_tensor)return tensor_from_list, tensor_from_np, zeros_tensor, ones_tensor, rand_tensor, eye_tensor# 2. Tensor 属性与类型转换
def tensor_properties_and_conversion(tensor_from_list, tensor_from_np):# 查看属性print("\nTensor属性:")print("形状:", tensor_from_list.shape)print("维度:", tensor_from_list.dim())print("元素总数:", tensor_from_list.numel())print("数据类型:", tensor_from_list.dtype)print("设备:", tensor_from_list.device)print("是否需要梯度:", tensor_from_list.requires_grad)# 类型转换float_tensor = tensor_from_list.float()  # 转为float32int_tensor = tensor_from_list.int()      # 转为int32long_tensor = tensor_from_list.long()    # 转为int64print("\n类型转换:")print("float32:", float_tensor.dtype)print("int32:", int_tensor.dtype)print("int64:", long_tensor.dtype)# 转为NumPy数组np_array = tensor_from_list.numpy()print("\n转为NumPy数组:\n", np_array)print("NumPy数组类型:", type(np_array))# 注意：共享内存tensor_from_np[0, 0] = 100print("修改Tensor后原NumPy数组:\n", np_array)# 3. Tensor 操作
def tensor_operations(rand_tensor):# 形状操作reshaped = rand_tensor.reshape(3, 2)  # 改变形状transposed = rand_tensor.t()          # 转置flattened = rand_tensor.flatten()     # 展平print("\n形状操作:")print("原形状:", rand_tensor.shape)print("reshape(3,2):", reshaped.shape)print("转置:", transposed.shape)print("展平:", flattened.shape)# 索引与切片print("\n索引与切片:")print("原始张量:\n", rand_tensor)print("第一行:", rand_tensor[0])print("第一列:", rand_tensor[:, 0])print("切片:\n", rand_tensor[0:1, 0:2])# 数学运算tensor1 = torch.randn(2, 3)tensor2 = torch.randn(2, 3)print("\n数学运算:")print("张量1:\n", tensor1)print("张量2:\n", tensor2)print("加法:\n", tensor1 + tensor2)print("减法:\n", tensor1 - tensor2)print("乘法:\n", tensor1 * tensor2)  # 元素乘法print("除法:\n", tensor1 / tensor2)print("矩阵乘法:\n", tensor1 @ tensor2.t())  # 矩阵乘法# 广播机制scalar = 2.0tensor_scalar_mul = tensor1 * scalar  # 标量广播print("\n标量乘法:\n", tensor_scalar_mul)vector = torch.randn(3)tensor_vector_add = tensor1 + vector  # 向量广播print("向量加法:\n", tensor_vector_add)# 4. 设备与梯度
def tensor_device_and_grad():# 设备移动tensor = torch.randn(2, 3)print("\n设备移动:")print("初始设备:", tensor.device)if torch.cuda.is_available():tensor_gpu = tensor.to('cuda')  # 移动到GPUprint("移动后设备:", tensor_gpu.device)tensor_cpu = tensor_gpu.cpu()  # 移动回CPUprint("移回CPU后设备:", tensor_cpu.device)# 梯度相关x = torch.tensor([2.0], requires_grad=True)y = x **2 + 3*x + 1print("\n梯度计算:")print("y =", y)y.backward()  # 反向传播print("x的梯度 (dy/dx):", x.grad)  # 应为2*x + 3 = 7if __name__ == "__main__":print("=== Tensor 创建 ===")tensors = create_tensors()print("\n=== Tensor 属性与类型转换 ===")tensor_properties_and_conversion(tensors[0], tensors[1])print("\n=== Tensor 操作 ===")tensor_operations(tensors[4])print("\n=== 设备与梯度 ===")tensor_device_and_grad()

Dataset 类

Dataset 类是 PyTorch 中用于封装数据的抽象类，自定义数据集需要继承该类并实现 __len__ 和 __getitem__ 方法。

自定义 Dataset 实现

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
from PIL import Image
import torchvision.transforms as transforms# 1. 基础自定义数据集
class SimpleDataset(Dataset):"""简单数据集示例"""def __init__(self, data_size=1000, input_dim=10, num_classes=5):"""参数:data_size: 数据集大小input_dim: 输入特征维度num_classes: 类别数量"""self.data_size = data_sizeself.input_dim = input_dimself.num_classes = num_classes# 生成随机数据self.data = torch.randn(data_size, input_dim)self.labels = torch.randint(0, num_classes, (data_size,))def __len__(self):"""返回数据集大小"""return self.data_sizedef __getitem__(self, idx):"""根据索引返回数据和标签"""return self.data[idx], self.labels[idx]# 2. 图像数据集
class ImageDataset(Dataset):"""图像数据集示例"""def __init__(self, image_dir, label_file, transform=None):"""参数:image_dir: 图像文件夹路径label_file: 标签文件路径transform: 图像预处理变换"""self.image_dir = image_dirself.transform = transform# 读取标签文件self.labels = {}with open(label_file, 'r') as f:for line in f:img_name, label = line.strip().split(',')self.labels[img_name] = int(label)# 获取图像文件名列表self.image_names = list(self.labels.keys())def __len__(self):return len(self.image_names)def __getitem__(self, idx):# 获取图像路径img_name = self.image_names[idx]img_path = os.path.join(self.image_dir, img_name)# 加载图像try:image = Image.open(img_path).convert('RGB')except Exception as e:print(f"加载图像 {img_name} 失败: {e}")# 返回一个随机图像作为替代image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))# 获取标签label = self.labels[img_name]# 应用变换if self.transform:image = self.transform(image)return image, label# 3. 带数据增强的数据集
class AugmentedDataset(Dataset):"""带数据增强的数据集"""def __init__(self, base_dataset, train=True):"""参数:base_dataset: 基础数据集train: 是否为训练模式（训练模式使用增强）"""self.base_dataset = base_datasetself.train = train# 定义训练和测试的变换self.train_transform = transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.RandomRotation(15),transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])self.test_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])])def __len__(self):return len(self.base_dataset)def __getitem__(self, idx):# 获取基础数据data, label = self.base_dataset[idx]# 应用相应的变换if self.train and isinstance(data, Image.Image):data = self.train_transform(data)elif not self.train and isinstance(data, Image.Image):data = self.test_transform(data)return data, label# 4. 数据集使用示例
def dataset_demo():# 1. 测试简单数据集simple_dataset = SimpleDataset(data_size=100, input_dim=10, num_classes=5)print("简单数据集大小:", len(simple_dataset))# 获取一个样本sample_data, sample_label = simple_dataset[0]print("样本数据形状:", sample_data.shape)print("样本标签:", sample_label)# 2. 测试图像数据集（这里使用模拟数据）# 注意：实际使用时需要提供真实的图像目录和标签文件class MockImageDataset(ImageDataset):"""模拟图像数据集，用于演示"""def __init__(self):super().__init__("", "")self.image_names = [f"image_{i}.jpg" for i in range(100)]self.labels = {name: i % 10 for i, name in enumerate(self.image_names)}def __getitem__(self, idx):img_name = self.image_names[idx]label = self.labels[img_name]# 创建随机图像image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))return image, labelimage_dataset = MockImageDataset()print("\n图像数据集大小:", len(image_dataset))# 获取一个样本img, label = image_dataset[0]print("图像模式:", img.mode)print("图像大小:", img.size)print("图像标签:", label)# 3. 测试带增强的数据集augmented_dataset = AugmentedDataset(image_dataset, train=True)img_augmented, label = augmented_dataset[0]print("\n增强后图像形状:", img_augmented.shape)print("增强后图像数据类型:", img_augmented.dtype)if __name__ == "__main__":dataset_demo()

DataLoader 类

DataLoader 用于将 Dataset 转换为可迭代的数据加载器，支持批量处理、打乱数据、多进程加载等功能。

DataLoader 用法与优化

import torch
from torch.utils.data import DataLoader, Dataset
import time
import numpy as np# 1. 创建一个大型数据集用于演示
class LargeDataset(Dataset):def __init__(self, size=10000, feature_dim=1000):self.size = sizeself.feature_dim = feature_dim# 生成大型数据集（实际中可能从文件读取）self.data = torch.randn(size, feature_dim)self.labels = torch.randint(0, 10, (size,))def __len__(self):return self.sizedef __getitem__(self, idx):# 模拟耗时的数据加载过程time.sleep(0.001)  # 模拟IO延迟return self.data[idx], self.labels[idx]# 2. DataLoader 基本用法
def dataloader_basics():dataset = LargeDataset(size=1000, feature_dim=100)# 基本DataLoaderbasic_loader = DataLoader(dataset,batch_size=32,shuffle=False,num_workers=0  # 单进程)# 计时start_time = time.time()for batch_idx, (data, labels) in enumerate(basic_loader):# 模拟训练处理if batch_idx % 10 == 0:print(f"批次 {batch_idx}, 数据形状: {data.shape}, 标签形状: {labels.shape}")basic_time = time.time() - start_timeprint(f"基本DataLoader耗时: {basic_time:.4f}秒")return dataset# 3. DataLoader 优化配置
def dataloader_optimization(dataset):# 优化的DataLoaderoptimized_loader = DataLoader(dataset,batch_size=32,shuffle=True,          # 打乱数据num_workers=4,         # 多进程加载pin_memory=True,       # 锁页内存prefetch_factor=2,     # 预取因子persistent_workers=True  # 保持工作进程)# 计时start_time = time.time()for batch_idx, (data, labels) in enumerate(optimized_loader):# 模拟训练处理passoptimized_time = time.time() - start_timeprint(f"优化DataLoader耗时: {optimized_time:.4f}秒")print(f"速度提升: {basic_time / optimized_time:.2f}x")# 不同批量大小的比较batch_sizes = [16, 32, 64, 128, 256]times = []for bs in batch_sizes:loader = DataLoader(dataset,batch_size=bs,num_workers=4,pin_memory=True)start = time.time()for data, labels in loader:passend = time.time()times.append(end - start)print(f"批量大小 {bs} 耗时: {end - start:.4f}秒")return batch_sizes, times# 4. 自定义collate_fn
def custom_collate_fn():"""自定义数据拼接函数，处理不同长度的数据"""class VariableLengthDataset(Dataset):def __init__(self, size=100):self.size = sizedef __len__(self):return self.sizedef __getitem__(self, idx):# 生成不同长度的序列length = np.random.randint(5, 20)data = torch.randn(length)label = torch.randint(0, 2, (1,)).item()return data, label# 自定义collate_fndef collate_fn(batch):"""处理不同长度的序列参数:batch: 包含元组(data, label)的列表"""# 分离数据和标签data_list, labels = zip(*batch)# 计算最大长度max_length = max(len(data) for data in data_list)# 填充序列到相同长度padded_data = []for data in data_list:pad_length = max_length - len(data)padded = torch.nn.functional.pad(data, (0, pad_length))padded_data.append(padded)# 堆叠数据和标签return torch.stack(padded_data), torch.tensor(labels)# 创建数据集和数据加载器dataset = VariableLengthDataset()loader = DataLoader(dataset,batch_size=10,collate_fn=collate_fn)# 测试data, labels = next(iter(loader))print("\n自定义collate_fn处理后:")print("数据形状:", data.shape)  # 应为(10, max_length)print("标签形状:", labels.shape)if __name__ == "__main__":print("=== DataLoader 基本用法 ===")dataset = dataloader_basics()print("\n=== DataLoader 优化配置 ===")batch_sizes, times = dataloader_optimization(dataset)print("\n=== 自定义collate_fn ===")custom_collate_fn()

完整数据处理流水线

结合 Tensor、Dataset 和 DataLoader，构建完整的数据处理流水线：

import torch
from torch.utils.data import Dataset, DataLoader, random_split
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image# 1. 构建完整的数据集类
class CompleteDataset(Dataset):def __init__(self, data_dir, split='train', transform=None):self.data_dir = data_dirself.split = splitself.transform = transform# 实际应用中，这里会加载真实数据# 为了演示，我们生成模拟数据self.num_samples = 1000 if split == 'train' else 200self.num_classes = 10def __len__(self):return self.num_samplesdef __getitem__(self, idx):# 生成模拟图像数据 (3通道, 224x224)image = Image.fromarray(np.uint8(np.random.randint(0, 256, (224, 224, 3))))# 生成随机标签label = np.random.randint(0, self.num_classes)# 应用变换if self.transform:image = self.transform(image)return image, label# 2. 构建数据处理流水线
def build_data_pipeline(data_dir, batch_size=32, val_split=0.2):# 定义数据变换train_transform = transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.RandomRotation(15),transforms.ColorJitter(brightness=0.2, contrast=0.2),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])])# 创建数据集train_dataset = CompleteDataset(data_dir, split='train',transform=train_transform)# 分割训练集和验证集val_size = int(len(train_dataset) * val_split)train_size = len(train_dataset) - val_sizetrain_dataset, val_dataset = random_split(train_dataset, [train_size, val_size])# 验证集使用验证变换val_dataset = CompleteDataset(data_dir, split='val',transform=val_transform)# 创建测试集test_dataset = CompleteDataset(data_dir, split='test',transform=val_transform)# 创建数据加载器train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=4,pin_memory=True,persistent_workers=True)val_loader = DataLoader(val_dataset,batch_size=batch_size,shuffle=False,num_workers=2,pin_memory=True)test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=2,pin_memory=True)return {'train': train_loader,'val': val_loader,'test': test_loader}# 3. 使用数据流水线进行训练
def train_with_data_pipeline(loaders, model, criterion, optimizer, num_epochs=10):device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')model.to(device)for epoch in range(num_epochs):print(f"\nEpoch {epoch+1}/{num_epochs}")print("-" * 10)# 训练阶段model.train()running_loss = 0.0running_corrects = 0for inputs, labels in loaders['train']:inputs = inputs.to(device)labels = labels.to(device)# 清零梯度optimizer.zero_grad()# 前向传播outputs = model(inputs)_, preds = torch.max(outputs, 1)loss = criterion(outputs, labels)# 反向传播和优化loss.backward()optimizer.step()# 统计running_loss += loss.item() * inputs.size(0)running_corrects += torch.sum(preds == labels.data)epoch_loss = running_loss / len(loaders['train'].dataset)epoch_acc = running_corrects.double() / len(loaders['train'].dataset)print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')# 验证阶段model.eval()running_loss_val = 0.0running_corrects_val = 0with torch.no_grad():for inputs, labels in loaders['val']:inputs = inputs.to(device)labels = labels.to(device)outputs = model(inputs)_, preds = torch.max(outputs, 1)loss = criterion(outputs, labels)running_loss_val += loss.item() * inputs.size(0)running_corrects_val += torch.sum(preds == labels.data)epoch_loss_val = running_loss_val / len(loaders['val'].dataset)epoch_acc_val = running_corrects_val.double() / len(loaders['val'].dataset)print(f'Val Loss: {epoch_loss_val:.4f} Acc: {epoch_acc_val:.4f}')return model# 4. 主函数
def main():# 构建数据流水线data_loaders = build_data_pipeline(data_dir='./data')# 显示一些样本inputs, labels = next(iter(data_loaders['train']))print(f"批量数据形状: {inputs.shape}")print(f"批量标签形状: {labels.shape}")# 创建简单模型class SimpleModel(torch.nn.Module):def __init__(self, num_classes=10):super(SimpleModel, self).__init__()self.features = torch.nn.Sequential(torch.nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1),torch.nn.ReLU(),torch.nn.MaxPool2d(kernel_size=2),torch.nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1),torch.nn.ReLU(),torch.nn.MaxPool2d(kernel_size=2),)self.classifier = torch.nn.Sequential(torch.nn.Flatten(),torch.nn.Linear(32 * 14 * 14, 128),torch.nn.ReLU(),torch.nn.Linear(128, num_classes))def forward(self, x):x = self.features(x)x = self.classifier(x)return xmodel = SimpleModel()# 定义损失函数和优化器criterion = torch.nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 训练模型model = train_with_data_pipeline(data_loaders, model, criterion, optimizer, num_epochs=3)if __name__ == '__main__':main()