devbox
Monodyee
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

PyTorch中torchvision库的详细介绍

作者：Monodyee | 2024-02-14 22:16:43

赞

踩

PyTorch中torchvision库的详细介绍

torchvision 是 PyTorch 生态系统中的一个关键库，专门为计算机视觉任务设计和优化。它提供了以下几个核心功能：

  1. 数据集：内置了多种广泛使用的图像和视频数据集，如 MNIST、CIFAR10/100、Fashion-MNIST、ImageNet、COCO 等，并且它们以 torch.utils.data.Dataset 的形式实现，方便与 PyTorch 数据加载器（DataLoader）集成。

  2. 数据预处理工具：通过 torchvision.transforms 模块提供了丰富的数据增强和预处理操作，包括但不限于裁剪、旋转、翻转、归一化、调整大小、颜色转换等，这些操作对于训练稳健的深度学习模型至关重要。

  3. 深度学习模型架构：在 torchvision.models 中封装了大量经典的预训练模型结构，例如 AlexNet、VGG、ResNet、Inception 系列、DenseNet、SqueezeNet 以及一些用于目标检测和语义分割任务的模型，用户可以直接加载这些模型进行迁移学习或者作为基础网络结构进行扩展。

  4. 实用工具：包含了一系列实用方法，比如将张量保存为图像文件、创建图像网格以便可视化多个样本等。

总之，torchvision 为基于 PyTorch 构建计算机视觉项目提供了极大的便利性，涵盖了从数据获取到模型构建及实验结果可视化等各个环节所需的功能。

1. 数据集

torchvision 是 PyTorch 的一个官方库，主要用于计算机视觉任务，它为开发者提供了一系列常用的数据集、模型架构以及图像转换工具。在 torchvision.datasets 子模块中，它包含了多个内置数据集，这些数据集可以直接用于训练和评估图像分类、对象检测、语义分割等多种视觉模型。以下是几个 torchvision 库中包含的常见数据集：

  1. MNIST：

    手写数字识别数据集，包含60,000个训练样本和10,000个测试样本，每个样本都是大小为28x28像素的单通道灰度图像，对应的标签是0-9的数字类别。

  2. CIFAR-10/100：

    • CIFAR-10 包含了60,000张32x32像素的彩色图像，分为10个类别，每类各有6000个样本（50,000用于训练，10,000用于测试）。
    • CIFAR-100 与 CIFAR-10 类似，但具有100个类别，每个类别有600张图片，因此对于细粒度分类更具挑战性。

  3. Fashion-MNIST：

    作为 MNIST 数据集的替代品，同样包含60,000训练样本和10,000测试样本，但是每个样本是一张28x28像素的时尚物品（如衬衫、裤子等）的灰度图像。

  4. ImageNet：

    虽然 torchvision 自身不直接提供 ImageNet 数据集的下载功能，但它提供了接口来加载已经下载好的 ILSVRC 2012 分类数据集（即通常所说的 ImageNet），该数据集包含超过1000类的物体类别，每类有数千张不同大小的RGB彩色图像。

  5. STL10：

    STL-10是一个小规模版本的ImageNet，有10个类别的100,000张未标记图像、5000张带标签的训练图像和8000张带标签的测试图像。

  6. COCO (Common Objects in Context)：

    COCO 数据集用于目标检测、分割和图像字幕等任务，包含大量标注的日常场景图片，每张图片可以包含多个目标及其边界框和分割掩模。

使用时，可以通过以下方式加载这些数据集：

Python
  1. 1import torch
  2. 2import torchvision
  3. 3from torchvision import datasets
  4. 4
  5. 5# 加载CIFAR-10数据集并进行基本处理
  6. 6transform = torchvision.transforms.Compose([...])  # 定义数据预处理操作
  7. 7dataset_train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
  8. 8dataset_test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
  9. 9
  10. 10# 使用DataLoader进一步将数据集转化为适合训练的批次
  11. 11dataloader_train = torch.utils.data.DataLoader(dataset_train, batch_size=..., shuffle=True)
  12. 12dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=...)

其中，root 参数指定了数据集存储的位置；train 参数确定是否加载训练集或测试集；download 参数设置为 True 则会自动从网上下载数据集；transform 参数允许对原始图像数据进行必要的预处理操作，例如归一化、裁剪、旋转等。

2. 数据预处理工具

torchvision 库中的数据预处理工具主要体现在 torchvision.transforms 模块，它提供了丰富的函数和类来对图像数据进行各种形式的转换和预处理。这些预处理操作在深度学习中是至关重要的，因为它们可以增强模型的泛化能力，并且将不同大小和格式的原始图像数据转化为神经网络能够接受的标准输入。

以下是一些常用的数据预处理方法：

  1. Resize：

    transforms.Resize(size, interpolation)：调整图像大小到指定尺寸。

  2. CenterCrop：

    transforms.CenterCrop(size)：从图像中心裁剪出一个给定大小的正方形区域。

  3. RandomCrop：

    transforms.RandomCrop(size, padding=None, pad_if_needed=False, fill=0, padding_mode='constant')：随机裁剪图像的一块区域。

  4. Normalize：

    transforms.Normalize(mean, std)：对图像像素值进行标准化处理，通常用于归一化RGB通道的均值和标准差。

  5. ToTensor：

    transforms.ToTensor()：将 PIL Image 或 numpy.ndarray 转换为 PyTorch 张量（从 0-255 的整数范围转换到 0-1 的浮点数范围）。

  6. ConvertImageDtype：

    transforms.ConvertImageDtype(dtype)：将图像转换为指定的数据类型。

  7. RandomHorizontalFlip：

    transforms.RandomHorizontalFlip(p=0.5)：以一定概率水平翻转图像。

  8. RandomVerticalFlip：

    transforms.RandomVerticalFlip(p=0.5)：以一定概率垂直翻转图像。

  9. ColorJitter：

    transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0)：随机改变图像的颜色属性，如亮度、对比度、饱和度和色调。

  10. Grayscale：

    transforms.Grayscale(num_output_channels=1)：将彩色图像转换为灰度图像。

  11. RandomRotation：

    transforms.RandomRotation(degrees, resample=False, expand=False, center=None, fill=None)：随机旋转图像。

为了方便使用，用户通常会组合多个预处理步骤，利用 transforms.Compose 类将其封装成一个预处理流水线：

Python
  1. 1from torchvision import transforms
  2. 2
  3. 3preprocess = transforms.Compose([
  4. 4    transforms.Resize(256),
  5. 5    transforms.CenterCrop(224),
  6. 6    transforms.ToTensor(),
  7. 7    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  8. 8])
  9. 9
  10. 10# 使用预处理流水线
  11. 11transformed_image = preprocess(image)

以上是对 torchvision.transforms 中一些关键数据预处理功能的概述，实际上该模块包含了更多丰富的方法，可以满足各种计算机视觉任务的需求。

3. 深度学习模型架构

torchvision.models 模块提供了大量预训练的深度学习模型，这些模型主要针对图像分类、对象检测和语义分割等计算机视觉任务。以下是一些常见的模型架构：

  1. 图像分类模型：

    • resnet18, resnet34, resnet50, resnet101, resnet152：基于残差网络（ResNet）架构，是目前最常用的深度神经网络之一，用于ImageNet数据集上的图像分类任务。
    • vgg16, vgg19：基于VGG（Visual Geometry Group）架构，特征提取能力强，但计算复杂度相对较高。
    • densenet121, densenet169, densenet201, densenet161：密集连接网络（DenseNet），通过密集块之间的稠密连接减少信息丢失并提升模型性能。
    • alexnet, squeezenet1_0, squeezenet1_1：AlexNet和SqueezeNet是较早的深度学习模型，前者在ILSVRC 2012竞赛中取得了突破性成果，后者以其轻量级结构著称。
    • googlenet, shufflenet_v2_x1_0, mobilenet_v2 等：为移动设备或资源受限环境设计的小型化网络。

  2. 对象检测模型：

    • torchvision不直接提供完整的预训练对象检测模型，但它包含了如 ssd, faster_rcnn 等检测模型的基本组件，用户可以利用 torchvision.ops 和 torchvision.models.detection 中的模块来构建自己的检测模型。

  3. 语义分割模型：

    • fcn_resnet50, deeplabv3_resnet50, lraspp_mobilenet_v3_large 等：全卷积网络（Fully Convolutional Networks, FCN）、DeepLabV3 和 MobileNetV3 架构为基础的语义分割模型，可用于像素级别的图像分类任务。

所有这些模型都支持加载预训练权重，并且能够作为基础结构进行迁移学习或微调以适应新的任务。例如，加载预训练的ResNet50模型进行图像分类任务可以通过以下方式实现：

Python
  1. 1import torchvision.models as models
  2. 2
  3. 3# 加载预训练模型
  4. 4model = models.resnet50(pretrained=True)
  5. 5
  6. 6# 将模型的最后一层替换为与新任务类别数匹配的线性层
  7. 7num_classes = len(new_dataset.classes)
  8. 8model.fc = nn.Linear(model.fc.in_features, num_classes)
  9. 9
  10. 10# 设定优化器并开始训练
  11. 11optimizer = torch.optim.Adam(model.parameters())

请注意，具体的模型列表可能会随着 torchvision 版本的更新而有所变化，因此建议查阅最新的官方文档获取详细信息。

4. 实用工具

torchvision 库除了提供数据集和预训练模型之外，还包含一些实用工具函数和类，这些工具在处理计算机视觉任务时非常有用。以下是一些关键的实用工具：

  1. 图像保存与读取：

    • torchvision.utils.save_image(tensor, filename, format=None)：将一个张量（通常是经过处理后的图像）保存为指定格式（如PNG、JPEG等）的图像文件。
    • 通过 PIL 或其他图像库读取图像后，可以使用 transforms.ToTensor() 将其转换为 PyTorch 张量。

  2. 图像显示：

    虽然 torchvision 自身不直接提供图像显示功能，但可以通过与外部库（如 matplotlib）结合来展示图像。例如，plt.imshow(torchvision.utils.make_grid(images)) 可以用来创建并显示一张由多个图像组成的网格。

  3. 图像拼接：

    torchvision.utils.make_grid(tensor, nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0)：将一组张量按照行列方式排列成一个大图像，常用于可视化多幅小图的结果。

  4. 视频处理：

    • torchvision.io.read_video(filename, start_pts=0, end_pts=float('inf'), pts_unit='sec', decoder_backend=None)：从视频文件中读取帧，并返回一个包含所有帧的 Tensor 列表。
    • torchvision.ops.video_reader.VideoReader(file_path, mode='video', backend=None)：提供了一个视频读取器对象，可用于逐帧读取视频。

  5. 图像元数据获取：

    对于某些加载的数据集（如COCO），torchvision 提供了方法来访问图像尺寸、标签以及其他元数据信息。

  6. 模型可视化工具：

    虽然不是严格意义上的“实用工具”，但 torchvision.models.utils 模块提供了用于生成模型结构图的方法，如 plot_model(model, show_shapes=False, to_file=None) 可以生成模型结构图（需要安装额外依赖如graphviz）。

以上列举的功能可以帮助开发者在进行计算机视觉任务时更好地管理和可视化数据及模型。同时，随着 torchvision 的不断更新和发展，可能会有更多实用工具加入其中。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Monodyee/article/detail/82104
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号