PyTorch 框架 CV 开发：从数据处理到模型部署与可视化

本文链接：https://blog.csdn.net/weixin_52603404/article/details/145204286

一、引言

在计算机视觉（CV）领域，PyTorch 作为一款强大且灵活的深度学习框架，为开发者提供了丰富的工具和资源，助力实现各类复杂的 CV 任务。从数据的处理与加载，到模型的训练、可视化以及最终的部署，PyTorch 构建了一套完整的技术栈。本文将依据给定的四篇文档，深入探讨 PyTorch 在 CV 开发中的进阶应用，涵盖自定义数据集类的创建、模型的训练与优化、训练过程的可视化以及模型的格式转换与推理，旨在为读者呈现一个全面且深入的 PyTorch CV 开发指南。

二、PyTorch 中的数据与数据集类

2.1 Pytorch 数据集类概述

PyTorch 中的数据处理是模型训练的基础，torch.util.data包提供了一系列用于数据加载和处理的工具。其中，Dataset是一个数据抽象类，它支持两种风格的数据表示：Map-style 和 Iterable-style。

Map-style 数据集：需实现__getitem__和__len__方法。__getitem__方法根据索引idx获取数据集中的样本，通过访问dataset[idx]即可获取对应的数据样本；__len__方法则返回数据集的大小。这种风格适用于可以通过索引快速访问的数据，如图片数据集，每个图片都有对应的索引，方便获取。
Iterable-style 数据集：通过迭代的方式获取数据，适用于数据无法通过索引直接访问的场景，例如从数据库中流式读取数据。

DataLoader则用于加载这两种风格的数据，它具有诸多强大的功能：支持自定义数据集加载，能够自动进行 batch 提取，实现单个或者多个数据加载，还可以自动进行内存固定，提高数据加载效率。其函数定义如下：

torch.utils.data.DataLoader(
    dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, 
    collate_fn=None,
    pin_memory=False, drop_last=False, timeout=0, 
    worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=2, persistent_workers=False
)

其中，batch_size指定每个批次的样本数量；shuffle决定是否在每个 epoch 对数据进行打乱；num_workers表示用于加载数据的子进程数量，增加该值可以加快数据加载速度，但会占用更多系统资源；pin_memory若设置为True，可将数据加载到锁页内存，加快数据从 CPU 到 GPU 的传输速度。

2.2 数据集预处理

在数据加载过程中，数据预处理是至关重要的环节。torchvision.transform提供了丰富的数据预处理工具，其中ToTensor可将数据转换为tensor数据，且取值范围为 0 - 1。例如，在图像数据集中，通常需要将图像转换为tensor格式，以便输入到模型中进行训练。通过ToTensor转换后，图像的像素值会被归一化到 0 - 1 的范围，这有助于模型的训练和收敛。

2.3 自定义图像数据集类代码演示

以自定义图像数据集类为例，假设我们有一个包含图像及其标注信息的数据集，标注信息可能包括图像中物体的位置、类别等。以下是一个简单的自定义数据集类的代码示例：

import torch
from torch.utils.data import Dataset

class CustomImageDataset(Dataset):
    def __init__(self, image_paths, annotations, transform=None):
        self.image_paths = image_paths
        self.annotations = annotations
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        annotation = self.annotations[idx]
        image = # 读取图像的代码
        if self.transform:
            image = self.transform(image)
        return image, annotation

在上述代码中，CustomImageDataset类继承自Dataset类。在__init__方法中，初始化图像路径列表、标注信息列表以及数据预处理转换操作。__len__方法返回数据集的大小，即图像的数量。__getitem__方法根据索引获取图像和对应的标注信息，并在有数据预处理操作时，对图像进行转换。通过这样的自定义数据集类，我们可以方便地加载和处理自己的图像数据集。

三、CNN 训练与测试

3.1 CUDA 支持

在模型训练过程中，CUDA 的支持可以极大地加速训练速度。PyTorch 提供了基于 CUDA 的训练方式，需要将数据和模型都迁移到 CUDA 设备上。具体步骤如下：

# 检查是否有可用的CUDA设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将数据迁移到CUDA设备
data = data.to(device)
# 将模型迁移到CUDA设备
model = model.to(device)

通过上述代码，首先检查系统中是否存在可用的 CUDA 设备，如果有，则将设备指定为cuda:0，否则使用 CPU。然后，将数据和模型分别迁移到指定的设备上，这样在训练过程中，计算将在 CUDA 设备上进行，从而提高训练效率。

3.2 定义损失与优化器

在训练模型时，需要定义损失函数和优化器。常见的优化器有 Adam 和 SGD。

Adam 优化器：Adam（Adaptive Moment Estimation）是一种自适应学习率的优化算法，它结合了 Adagrad 和 RMSProp 算法的优点，能够根据参数的梯度自适应地调整学习率。在训练过程中，Adam 优化器可以更快地收敛到最优解，尤其适用于数据量较大、参数较多的模型。
SGD 优化器：SGD（Stochastic Gradient Descent）即随机梯度下降法，是一种简单且常用的优化算法。它通过计算每个样本的梯度来更新参数，具有计算简单、易于实现的特点。然而，SGD 在训练过程中可能会出现收敛速度慢、容易陷入局部最优等问题。相比之下，Adam 优化器在很多情况下能够更快地找到较优的解。

在实际应用中，需要根据模型的特点和数据的情况选择合适的优化器。例如，对于一些复杂的模型，Adam 优化器可能更适合；而对于简单的模型，SGD 优化器可能已经足够。

3.3 训练与预测代码演示

以下是一个简单的 CNN 训练与预测的代码示例：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 定义模型结构

    def forward(self, x):
        # 前向传播

# 创建模型实例
model = CNN()
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 将数据和标签迁移到CUDA设备
        data, target = data.to(device), target.to(device)
        # 梯度清零
        optimizer.zero_grad()
        # 前向传播
        output = model(data)
        # 计算损失
        loss = criterion(output, target)
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()

# 预测
with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        # 对输出进行处理，得到预测结果

在上述代码中，首先定义了一个简单的 CNN 模型，然后创建模型实例。接着，定义了交叉熵损失函数CrossEntropyLoss和 Adam 优化器，并设置学习率为 0.001。在训练过程中，遍历训练数据加载器，将数据和标签迁移到 CUDA 设备上，进行前向传播、计算损失、反向传播和参数更新。在预测阶段，遍历测试数据加载器，将数据迁移到 CUDA 设备上，通过模型进行前向传播，得到预测结果。

四、Pytorch 训练可视化

4.1 Tensorboard 可视化

Tensorboard 最初来自 TensorFlow，PyTorch 通过torch.utils.tensorboard模块支持 Tensorboard 可视化。使用 Tensorboard 可以方便地可视化模型的训练过程，包括损失值的变化、准确率的变化等。

安装与启动：首先需要安装 Tensorboard，通过pip install tensorboard命令即可完成安装。启动 Tensorboard 的命令为tensorboard --logdir=D:/log，其中D:/log是日志文件的存储路径。启动后，可以通过浏览器访问http://localhost:6006查看可视化结果。
标量与图像保存：在训练过程中，可以使用SummaryWriter将标量数据（如损失值、准确率等）和图像数据保存到日志文件中，以便在 Tensorboard 中进行可视化。以下是一个简单的示例：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs')
for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 训练过程
        loss = # 计算损失
        writer.add_scalar('Training Loss', loss, epoch * len(train_loader) + batch_idx)
writer.close()

在上述代码中，创建了一个SummaryWriter实例，并指定日志文件的存储路径为runs。在训练过程中，使用add_scalar方法将训练损失值保存到日志文件中，其中Training Loss是标量的名称，loss是损失值，epoch * len(train_loader) + batch_idx是全局步数。通过这样的方式，在 Tensorboard 中可以直观地看到训练损失值随时间的变化情况。

4.2 Visdom 可视化

Visdom 是另一种用于可视化的工具，其官方地址为https://github.com/noagarcia/visdomtutorial。Visdom 可以实时可视化模型的训练过程，支持绘制各种图表，如折线图、柱状图等。使用 Visdom 时，需要先安装 Visdom 库，然后在代码中导入并使用相关的函数进行可视化。以下是一个简单的 Visdom 可视化示例。

import visdom
import time

vis = visdom.Visdom()
win = vis.line(X=torch.zeros((1,)), Y=torch.zeros((1,)), opts=dict(xlabel='Iteration', ylabel='Loss', title='Training Loss'))

for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 训练过程
        loss = # 计算损失
        vis.line(X=torch.ones((1,)) * (epoch * len(train_loader) + batch_idx), Y=torch.Tensor([loss]), win=win, update='append')
        time.sleep(0.1)

在上述代码中，首先创建了一个Visdom实例vis，然后使用vis.line方法创建了一个折线图窗口win，用于显示训练损失值。在训练过程中，使用vis.line方法更新折线图，将当前的迭代次数和损失值添加到折线图中。通过这样的方式，可以实时观察训练损失值的变化情况。

五、ONNX 格式模型导出与推理

5.1 ONNX 格式说明

ONNX（Open Neural Network Exchange）是一种用于表示深度学习模型的开放标准格式。它允许在不同的深度学习框架之间进行模型的转换和共享。ONNX 的官方代码库位于https://github.com/onnx/onnx。与 ONNX 相关的还有 ONNX Runtime，它是一个高性能的推理引擎，用于运行 ONNX 格式的模型。此外，ONNX 还可以与 OpenVINO、OpenCV 等工具结合使用，实现模型的高效部署和推理，如果使用Pytorch框架生成的是pt/pth的权重文件，而模型的保存无法部署到其他的框架中，因此我们需要把pth/pt文件部署到其他框架上的时候，就需要将其转换成onnx的格式。

5.2 ONNX 格式模型转换

PyTorch 支持将pt或pth格式的模型转换为 ONNX 格式。在转换过程中，需要编写转换脚本，同时还可以查看模型的输入与输出、格式以及预处理等信息。以下是一个简单的模型转换示例，该项目是人类表情识别，使用Resnet网络训练模型，将模型导出为onnx格式：

import torch
import torchvision.models as models

def export_onnx(model, file_name="face_emotions_model.onnx", device='cuda'):
    """将模型导出为ONNX格式"""
    model.eval()  # 设置模型为评估模式
    dummy_input = torch.randn(1, 3, 224, 224).to(device)  # 假设输入是3通道的224x224图像
    model = model.to(device)  # 确保模型在同一设备上
    torch.onnx.export(model, dummy_input, file_name, verbose=True)
    print(f"ONNX模型已导出为 {file_name}")

if __name__ == "__main__":
    model = EmotionsResNet()
    print(model)

    if train_on_gpu:
        model.cuda()

    ds = EmotionDataset('')
    num_train_samples = ds.num_of_samples()
    bs = 16
    dataloader = DataLoader(ds, batch_size=bs, shuffle=True)

    num_epochs = 15
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
    model.train()

    cross_loss = torch.nn.CrossEntropyLoss()
    index = 0
    for epoch in range(num_epochs):
        train_loss = 0.0

        for i_batch, sample_batched in enumerate(dataloader):
            images_batch, emotion_batch = sample_batched['image'], sample_batched['emotion']

            if train_on_gpu:
                images_batch, emotion_batch = images_batch.cuda(), emotion_batch.cuda()

            optimizer.zero_grad()

            m_emotion_out = model(images_batch)
            emotion_batch = emotion_batch.long()

            loss = cross_loss(m_emotion_out, emotion_batch)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            if index % 100 == 0:
                print('step: {} \tTraining Loss: {:.6f}'.format(index, loss.item()))
            index += 1

        train_loss = train_loss / num_train_samples
        print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch, train_loss))

    model.eval()
    torch.save(model, 'face_emotions_model.pt')
    print("模型已保存为 face_emotions_model.pt")

    # 导出为ONNX格式
    export_onnx(model)

在上述代码中，首先创建了一个预训练的 ResNet-18 模型，并将其设置为评估模式。然后，创建了一个随机输入张量dummy_input，其形状为(1, 3, 224, 224)，表示一个包含 1 张 3 通道、大小为 224x224 的图像。最后，使用torch.onnx.export方法将模型导出为 ONNX 格式，resnet18.onnx是输出的 ONNX 模型文件名，verbose=True表示在转换过程中输出详细信息。

5.3 部署与预测代码演示

在将模型转换为 ONNX 格式后，可以使用 ONNX Runtime 或 OpenCV 等工具进行模型的部署和预测。以下是使用 ONNX Runtime 进行预测的示例：

import onnxruntime as ort
import numpy as np

# 创建ONNX Runtime推理会话
ort_session = ort.InferenceSession("resnet18.onnx")

# 加载图像并进行预处理
image = # 加载图像的代码
image = # 预处理图像的代码

# 进行推理
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name
result = ort_session.run([output_name], {input_name: image})

在上述代码中，首先创建了一个 ONNX Runtime 推理会话ort_session，并加载 ONNX 模型resnet18.onnx。然后，加载并预处理图像，将预处理后的图像作为输入数据。最后，使用ort_session.run方法进行推理，get_inputs和get_outputs方法用于获取模型的输入和输出节点信息，通过指定输入和输出节点的名称，得到推理结果。

六、小结

本文介绍了 PyTorch 在 CV 开发中的进阶应用，从数据的处理与加载，到模型的训练、可视化以及最终的部署。通过自定义数据集类，我们可以灵活地加载和处理各种数据；在模型训练方面，利用 CUDA 加速、选择合适的损失函数和优化器，能够提高模型的训练效率和性能；训练可视化工具如 Tensorboard 和 Visdom，让我们可以直观地了解模型的训练过程，便于进行模型的调优；最后，将模型转换为 ONNX 格式，为模型的跨框架部署和推理提供了便利。在实际的 CV 项目中，开发者可以根据具体的需求和场景，综合运用这些技术，实现高效、准确的计算机视觉任务。随着深度学习技术的不断发展，PyTorch 也在持续更新和完善，为 CV 开发者提供更多强大的工具和功能，助力推动计算机视觉领域的进一步发展，感谢大家的观看(●ˇ∀ˇ●)。