torchvision 教程

news2024/9/20 0:55:51

PyTorch torchvision 教程

torchvision 是 PyTorch 的一个子库，专为计算机视觉任务设计，提供了常用的数据集、预训练模型、以及图像转换和处理的工具。本文将介绍如何使用 torchvision 中的功能来加载数据集、预处理数据、使用预训练模型以及进行图像增强。

1. 安装 torchvision

首先，你需要安装 torchvision 库。可以使用 pip 安装：

pip install torchvision

2. torchvision 的主要组件

torchvision 的主要组件有：

  • torchvision.datasets：提供常用的数据集，例如 MNIST、CIFAR-10、ImageNet 等。
  • torchvision.transforms：用于图像的预处理和数据增强。
  • torchvision.models：提供预训练的深度学习模型。
  • torchvision.io：用于读取和写入图像、视频等数据。

3. 使用 torchvision.datasets 加载数据集

torchvision 提供了许多流行的数据集，可以直接从 torchvision.datasets 中加载。你可以加载数据集，并使用 DataLoader 迭代数据。

3.1 加载 MNIST 数据集

MNIST 是一个包含手写数字的经典数据集，常用于图像分类任务。

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 定义数据转换 (如归一化)
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 下载并加载 MNIST 数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 使用 DataLoader 加载数据集
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 打印样本
for images, labels in train_loader:
    print(f"Image batch shape: {images.size()}")
    print(f"Labels batch shape: {labels.size()}")
    break
3.2 加载 CIFAR-10 数据集

CIFAR-10 是另一个常用的数据集，包含 10 类自然图片。

# 加载 CIFAR-10 数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 使用 DataLoader
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

4. torchvision.transforms 图像预处理与增强

torchvision.transforms 提供了许多常用的图像预处理和增强方法，例如缩放、裁剪、旋转、翻转等。

4.1 基本预处理操作
transform = transforms.Compose([
    transforms.Resize((32, 32)),            # 调整图像大小
    transforms.RandomHorizontalFlip(),      # 随机水平翻转
    transforms.ToTensor(),                  # 转换为 PyTorch 张量
    transforms.Normalize((0.5,), (0.5,))    # 标准化
])

train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
4.2 常见的 transforms 操作
  • transforms.Resize(size)：调整图像大小为 size。
  • transforms.CenterCrop(size)：从图像中心裁剪大小为 size 的部分。
  • transforms.RandomCrop(size)：随机裁剪图像。
  • transforms.RandomHorizontalFlip()：随机水平翻转图像。
  • transforms.ColorJitter()：随机更改图像的亮度、对比度和饱和度。
  • transforms.ToTensor()：将 PIL 图像或 NumPy 数组转换为 PyTorch 张量。
  • transforms.Normalize(mean, std)：标准化图像数据。

5. 使用 torchvision.models 的预训练模型

torchvision.models 提供了多种预训练模型，例如 ResNet、VGG、AlexNet 等，这些模型在 ImageNet 数据集上进行了预训练。

5.1 加载预训练模型

你可以加载一个预训练的 ResNet 模型并在新任务上进行微调。

import torchvision.models as models

# 加载预训练的 ResNet18 模型
model = models.resnet18(pretrained=True)

# 查看模型架构
print(model)
5.2 微调预训练模型

如果你想要微调预训练模型（例如用于 CIFAR-10 数据集），你可以冻结预训练模型的部分参数，并修改最后一层以适应新的任务。

# 冻结所有层的参数
for param in model.parameters():
    param.requires_grad = False

# 修改最后一层以适应 CIFAR-10 (10 类分类任务)
model.fc = torch.nn.Linear(512, 10)

# 将模型移动到 GPU（如果有）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
5.3 训练模型
# 训练模型
for epoch in range(2):
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        # 清零梯度
        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播与优化
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')
5.4 测试模型
# 测试模型性能
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total}%')

6. torchvision.io 读取和保存图像

torchvision.io 提供了方便的图像读取和保存功能。

6.1 读取图像
import torchvision.io as io

# 读取图像
img = io.read_image('image.jpg')  # 读取为张量

# 显示张量信息
print(img.size())
6.2 保存图像
# 保存张量为图像文件
io.write_jpeg(img, 'output_image.jpg')

7. 完整示例

以下是一个使用 torchvision 加载数据集、进行数据增强、使用预训练模型微调并进行训练的完整示例：

import torch
import torchvision
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim

# 数据增强与预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载 CIFAR-10 数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 加载预训练的 ResNet18 模型并修改最后一层
model = models.resnet18(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
model.fc = nn.Linear(512, 10)

# 设备设置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(2):
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs

 = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')

# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total}%')

8. 总结

torchvision 是 PyTorch 中处理计算机视觉任务的重要工具，它为常用的数据集、模型、数据处理和增强提供了便利的接口。通过本教程，你可以学习如何使用 torchvision 加载数据集、应用图像预处理、使用预训练模型进行微调，并训练模型来解决实际的计算机视觉任务。

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.coloradmin.cn/o/2147450.html

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈，一经查实，立即删除！

相关文章

TIOBE 编程指数 9 月排行榜公布 VB.Net第七

TIOBE 编程指数 9 月排行榜公布 VB.Net第七

原文地址：百度安全验证 IT之家 9 月 8 日消息，TIOBE 编程社区指数是一个衡量编程语言受欢迎程度的指标，评判的依据来自世界范围内的工程师、课程、供应商及搜索引擎，今天 TIOBE 官网公布了 2024 年 9 月的编程语言排行榜&#xf…
阅读更多...
Element走马灯组件循环播放两个页面是方向不一致

Element走马灯组件循环播放两个页面是方向不一致

摘要：使用Carousel 走马灯循环播放同一类型的图片、文字等内容，会在循环内容为两组是出现下图 [1]中的现象。本文记录下如何解决 之前项目遇到过一次这个问题，由于indicator-position 指示器不用显示，则判断内容长度为2时&#xf…
阅读更多...
恶意Bot流量识别分析实践

恶意Bot流量识别分析实践

1、摘要 随着互联网的发展，自动化工具和脚本（Bots）的使用越来越普遍。虽然一些善意 Bots 对于网站的正常运行和数据采集至关重要，但恶意 Bots 可能会对网站带来负面影响，如爬取敏感信息、恶意注册、刷流量等。因此&am…
阅读更多...
saltstack配置管理

saltstack配置管理

一、saltstack的SSH工作模式 一、salt-ssh介绍 salt-ssh 是 0.17.0 新引入的一个功能，不需要minion对客户端进行管理，也不需要master。salt-ssh 支持salt大部分的功能：如 grains、modules、state 等salt-ssh 没有使用ZeroMQ的通信架构&#…
阅读更多...
向日葵好用吗？4款稳定的远程控制软件推荐。

向日葵好用吗？4款稳定的远程控制软件推荐。

远程控制技术现在已经被应用于很多个领域，像企业办公，远程协助，智能家居，工业控制等等。我们常常会用到的时前两种。而实现远程控制的方式也有多种，但是最方便高效的还是使用第三方软件。我最常使用的是向日葵&#xf…
阅读更多...
算法.图论-并查集上

算法.图论-并查集上

文章目录 1. 并查集介绍2. 并查集的实现2.1 实现逻辑2.2 isSameSet方法2.3 union方法(小挂大优化)2.4 find方法(路径压缩优化) 3. 并查集模板 1. 并查集介绍 定义： 并查集是一种树型的数据结构，用于处理一些不相交集合的合并及查询问题（即所…
阅读更多...
游戏如何对抗定制挂

游戏如何对抗定制挂

近年来，游戏安全对抗强度相比以往更加激烈，具体表现在“定制挂”趋势显著。在近期收集的近万款外挂样本中，定制挂约占比78%，常见的内存修改器、变速器等通用作弊手段占比正在下降。 所谓定制挂，是指针对某款游戏单独开…
阅读更多...
SPI接口通信协议浅谈成都自动化开发

SPI接口通信协议浅谈成都自动化开发

沙鸥-成都 1 什么是SPI SPI是串口外设接口的缩写，是一种高速的、全双工、同步的通信协议，是微处理器与外围IC之间常用的一种通讯方式。 SPI是主从式的通信协议，可以一主机一从机通信，也可以一主机多从机通信。 2 SPI的优缺点 SPI接…
阅读更多...
【Java版】云HIS系统源码

【Java版】云HIS系统源码

云HIS系统介绍 云HIS系统是一款满足基层医疗机构各类业务需要的健康云产品。该产品能帮助基层医疗机构完成日常各类业务，提供病患挂号支持、病患问诊、电子病历、开药发药、会员管理、统计查询、医生站和护士站等一系列常规功能，还能与公卫、PACS等各类…
阅读更多...
【STM32 HAL库】OLED显示模块

【STM32 HAL库】OLED显示模块

【STM32 HAL库】OLED显示模块 前言理论OLED基本参数OLED基本驱动原理OLED坐标轴 应用CubeMx配置底层函数代码高层封装函数printf显示函数 前言 本文为笔者学习 OLED 的总结，基于keysking的视频内容，如有错误，欢迎指正 理论 OLED基本参数 …
阅读更多...
基于224G的超高速以太网端口1.6Tbps 1600G真的来了~

基于224G的超高速以太网端口1.6Tbps 1600G真的来了~

基于224G PAM4 SerDes的1.6T(更激进些的是3.2T） 受AI智能算中心的驱动，基于4x112G的400G光模块，和基于8x112G的800G的光模块已经很成熟了，标志就是大家都在降本增效，考虑干掉功耗的DSP，...另外一个标志就是…
阅读更多...
关于支持向量机的一份介绍

关于支持向量机的一份介绍

在这篇文章中，我将介绍与支持向量机有关的东西，我们知道支持向量机主要分两类，就是线性支持向量机和核支持向量机这两种（当然还有其他的，如多类支持向量机、 Nu-Support Vector Regression等），因…
阅读更多...
AIGC生图基础知识

AIGC生图基础知识

一、引言 AIGC，即AI-Generated Content，是一种利用大型预训练模型如生成对抗网络（GAN）、扩散网络（Diffusion）和语言大模型（Transformer）等人工智能技术，通过对大量数据进…
阅读更多...
Gradio 自定义组件

Gradio 自定义组件

如何使用 Gradio 自定义组件，Gradio 前端使用 Svelte，后端使用的 Python。如何自定义一个组件呢？Gadio 提供了类似于脚手架的命令，可以生成需要开发组件的前后和后端代码。 创建组件 运行如下命令，gradio 会自动生成…
阅读更多...
OBC充电机测试的步骤和规范

OBC充电机测试的步骤和规范

一、测试前准备 1. 确认测试环境：确保测试环境的温度、湿度等条件符合设备的工作要求。 2. 检查设备：检查OBC充电机是否完好无损，电源线、插头等是否连接良好，显示屏是否正常显示。 3. 准备工具：准备好电压表、电流…
阅读更多...
ubuntu20.04安装cudnn

ubuntu20.04安装cudnn

先登入账号 网址：https://developer.nvidia.com/cudnn 选择ubuntu20.04 x86_64（Deb） 在下载好文件的文件夹下打开终端 sudo apt-get install zlib1gsudo dpkg -i cudnn-local-repo-${distro}-8.x.x.x_1.0-1_amd64.debsudo cp /var/cudnn-lo…
阅读更多...
大数据时代：历史、发展与未来

大数据时代：历史、发展与未来

文章目录 引言1980年：大数据的先声2006年：云计算与大数据的诞生2008年：大数据的科学探索2009年：大数据成为行业热词2011年：大数据的商业价值2013年：世界大数据元年结语 引言 在信息技术飞速发展的今天&…
阅读更多...
钢铁焦化水泥超低排的原因

钢铁焦化水泥超低排的原因

钢铁、焦化和水泥行业实施超低排放的原因，朗观视觉小编建议大家可以从环境保护、产业升级、政策推动以及企业可持续发展等多个方面进行分析。 一、环境保护需求 空气质量改善：钢铁、焦化和水泥行业是传统的高污染行业，其排放的颗粒物、二氧化…
阅读更多...
表格HTML

表格HTML

//test.html <!DOCTYPE html> <html lang"en"> <head> <meta charset"UTF-8"> <meta name"viewport" content"widthdevice-width, initial-scale1.0"> <title>表格与CSS分开示例</tit…
阅读更多...
【STL】string 基础，应用与操作

【STL】string 基础，应用与操作

string 1.string相关介绍 STL&#xff08;标准模板库&#xff09;中的string容器是C标准库提供的用于处理和操作字符串的类&#xff0c;位于头文件中。std::string提供了比传统的C风格字符串&#xff08;字符数组&#xff09;更方便和安全的功能&#xff0c;具有动态内存管理…
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023