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

PyTorch数据处理工具箱详解｜深入理解torchvision与torch.utils.data

news 2025/8/19 13:47:55

在深度学习的旅程中，数据处理是构建模型前不可或缺的一环。PyTorch 提供了一系列高效、灵活的数据处理工具，帮助开发者更便捷地完成数据装载、预处理、增强等任务。本文将围绕 PyTorch 中的核心数据处理工具 torch.utils.data 与 torchvision 展开详细介绍，并帮助读者理解它们之间的关系和使用场景。

一、核心数据处理引擎：torch.utils.data

位于图4-1左侧的是 PyTorch 提供的基础数据处理模块 torch.utils.data，它为数据集的定义、迭代、采样等提供了一系列类和函数。主要包括以下四个核心类：

1. Dataset（数据集抽象基类）

Dataset 是一个抽象类，所有自定义数据集都应继承此类。
需要实现以下两个方法：
- __getitem__(self, index)：根据索引返回单个样本；
- __len__(self)：返回数据集的总样本数。
作用：定义如何访问单个样本，是构建数据集的基础。

2. DataLoader（数据加载器）

DataLoader 是一个迭代器，用于按批次（batch）加载数据。
支持功能：
- 批量读取（batching）
- 数据打乱（shuffle）
- 并行加载（num_workers）
作用：将原始数据封装为可批量读取的数据流，是训练过程中的“数据管道”。

3. random_split（数据集划分工具）

可将一个数据集随机拆分为多个子集，如训练集、验证集和测试集。
保证子集之间无交集，适用于数据分割、交叉验证等场景。

示例：

train_dataset, val_dataset = random_split(full_dataset, [50000, 10000])

4. Sampler（采样器）

Sampler 是一系列采样策略类，控制数据的读取顺序。
常见采样器包括：
- SequentialSampler：顺序采样
- RandomSampler：随机采样
- SubsetRandomSampler：从子集中随机采样
- WeightedRandomSampler：带权重的随机采样
作用：在 DataLoader 中自定义采样逻辑，提升训练灵活性。

二、视觉处理工具箱：torchvision

中间部分介绍的是 torchvision，作为 PyTorch 的视觉扩展库，它独立于 PyTorch 主库，需通过以下命令单独安装：

pip install torchvision
或使用 conda 安装
conda install torchvision

torchvision 主要包含四大类功能模块，分别用于数据集处理、模型调用、图像预处理和图像操作。

1. datasets（常用视觉数据集）

提供了多个标准数据集接口，如：
- MNIST（手写数字识别）
- CIFAR-10 / CIFAR-100（彩色图像分类）
- ImageNet（大规模图像分类）
- COCO（目标检测与图像描述）
所有数据集都继承自 torch.utils.data.Dataset，可无缝接入 DataLoader。
优势：一键加载、统一接口、节省开发时间。

2. models（经典模型与预训练网络）

包含大量经典神经网络结构，如：
- AlexNet、VGG、ResNet、Inception 等
支持加载预训练模型（设置 pretrained=True），便于迁移学习。

示例：

import torchvision.models as models 
model = models.resnet18(pretrained=True)

3. transforms（图像变换操作）

提供对图像进行预处理和增强的功能。
支持的操作类型包括：
- 对 PIL 图像的操作（如 Resize、Crop、Normalize）
- 对 Tensor 的操作（如 ToTensor）

示例：

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])
])