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

AI开发-三方库-torch-torchvision

news 2026/2/9 0:44:07

1 需求

数据集：torchvision.datasets
- torchvision.datasets.MNIST
数据变换：torchvision.transforms
- torchvision.transforms.Compose
- torchvision.transforms.ToTensor
- torchvision.transforms.Normalize
模型：torchvision.models
可视化工具：torchvision.utils

2.1 torchvision.datasets

torchvision.datasets.MNIST

CLASS torchvision.datasets.MNIST(root: Union[str, Path], train: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, download: bool = False)

root
train
transform
download

MNIST — Torchvision 0.20 documentation

2.2 torchvision.transforms

torchvision.transforms
torchvision.transforms.Compose
torchvision.transforms.ToTensor
torchvision.transforms.Normalize

torchvision.transforms

torchvision.transforms — Torchvision 0.11.0 documentation

Transforming and augmenting images — Torchvision 0.20 documentation

CLASS torchvision.transforms.Compose(transforms)

Compose — Torchvision 0.20 documentation

CLASS torchvision.transforms.ToTensor

ToTensor — Torchvision 0.20 documentation

CLASS torchvision.transforms.Normalize(mean, std, inplace=False)

Normalize — Torchvision 0.20 documentation

3 示例

import torchvision.transforms as transforms
from PIL import Imagetransforms = transforms.Compose([transforms.ToTensor()
])image = Image.open("./Data/1.png")tensor_image = transforms(image)print(tensor_image)

torchvision 是一个基于 PyTorch 的用于计算机视觉任务的库，它提供了诸多便捷的工具、数据集和模型等资源，极大地简化了计算机视觉相关深度学习项目的开发流程。以下是对它的详细介绍：

3.1 数据集（Datasets）

torchvision.datasets 模块包含了许多常见的计算机视觉数据集的加载器，比如：
- MNIST：这是一个手写数字识别数据集，包含了大量的手写数字图像（0 - 9）以及对应的标签，常用于图像分类任务的入门练习和研究。
- CIFAR - 10：由 10 个不同类别的 60000 张彩色小图像组成，每类有 6000 张图像，图像尺寸较小（32×32 像素），同样广泛应用于图像分类任务的实验和模型评估等。
- CIFAR - 100：与 CIFAR - 10 类似，但包含 100 个不同的类别，图像数量和尺寸与 CIFAR - 10 相同，可用于更具挑战性的分类任务。
- ImageNet：一个非常大规模的图像数据集，涵盖了众多的类别（超过 1000 个）和大量的图像，是许多先进的计算机视觉模型训练和评估的重要基础数据集，不过其完整数据集体积庞大，使用时可能需要根据具体需求进行部分下载或采用其预训练模型。

这些数据集加载器使得获取和预处理这些常用数据集变得十分容易，只需简单设置参数（如是否下载、数据集根目录、数据变换等）就可以将数据集加载到内存并转换为适合模型训练的格式（通常是张量格式）。

3.2 数据变换（Transforms）

torchvision.transforms 模块提供了一系列用于对图像进行变换的操作，主要用于数据增强和数据预处理，具体如下：
- 几何变换：包括裁剪（如中心裁剪 CenterCrop、随机裁剪 RandomCrop）、翻转（如水平翻转 RandomHorizontalFlip、垂直翻转 RandomVerticalFlip）、旋转（RandomRotation）等操作，这些操作可以改变图像的几何形状，有助于模型学习到图像不同位置和姿态下的特征，同时也是数据增强的重要手段。
- 颜色变换：例如通过 ColorJitter 可以对图像的亮度、对比度、饱和度和色调进行随机调整，让模型能够适应不同颜色状态下的图像，提高模型的鲁棒性。
- 归一化（Normalize）：对图像的像素值进行归一化处理，通常是将像素值转换为均值为 0、标准差为 1 的分布，这样可以使数据分布更加符合模型训练的假设，有助于模型更快地收敛，避免梯度消失或梯度爆炸等问题。
- 类型转换（ToTensor）：将图像数据从常见的格式（如 PIL 图像或者 numpy 数组）转换为 torch 中的张量格式，并且在转换过程中会自动将像素值范围从 [0, 255] 转换为 [0, 1]（对于单通道图像）或者 [0, 1] 区间内的浮点数张量（对于多通道图像），这是使用 torchvision 数据集和模型的基本前提，因为 torch 中的模型需要以张量作为输入。

可以通过 torchvision.transforms.Compose 将多个变换操作组合在一起，形成一个完整的数据预处理流程。

3.3 模型（Models）

torchvision.models 模块提供了许多预训练好的计算机视觉模型，这些模型都是基于 PyTorch 构建的，常见的有：
- AlexNet：早期具有开创性意义的深度学习模型之一，在图像分类任务上取得了显著的成果，为后续模型的发展奠定了基础。
- VGG 系列：如 VGG16、VGG19 等，具有较深的网络结构，通过堆叠多个小卷积核的卷积层来实现较好的图像特征提取效果，在图像分类等任务中表现出色。
- ResNet 系列：包括 ResNet18、ResNet34、ResNet51 等，引入了残差连接的概念，有效解决了随着网络深度增加而出现的梯度消失或梯度爆炸问题，在图像分类、目标检测等多个计算机视觉任务中广泛应用。
- DenseNet 系列：采用了密集连接的方式，使得网络中的每个层都与前面的所有层直接相连，进一步提高了特征提取的效率和效果，在一些计算机视觉任务中也有很好的表现。

这些预训练模型可以直接用于预测任务，或者在其基础上进行微调（fine - tuning）以适应特定的视觉任务，如目标分类、目标检测等。通过简单设置参数就可以加载这些模型，并利用其已有的知识（即预训练的参数）来加速新任务的学习过程。

3.4 可视化工具（Utils）

torchvision.utils 模块提供了一些工具来帮助可视化图像数据和模型输出。例如：
- make_grid：可以将一批图像整理成一个网格形式的张量，便于在训练过程中查看图像批次的情况，了解数据的分布和变化。
- save_image：用于保存可视化后的图像，方便记录和分析训练过程中的各种图像相关结果。

这些可视化工具对于理解模型的行为和性能很有帮助，能够让开发者更直观地观察到模型在训练过程中的表现以及数据的处理情况。

总的来说，torchvision 是 PyTorch 生态系统中不可或缺的一部分，为计算机视觉领域的深度学习项目提供了全方位的支持。