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

Python 【图像分类】之 PyTorch 进行猫狗分类功能的实现(Swanlab训练可视化/ Gradio 实现猫狗分类 Demo)

news 来源：原创 2025/5/8 19:54:59

Python 【图像分类】之 PyTorch 进行猫狗分类功能的实现(Swanlab训练可视化/ Gradio 实现猫狗分类 Demo)

一、简单介绍

二、PyTorch

三、CNN

1、神经网络

2、卷积神经网络

四、ResNet50

五、Pytorch 实现猫狗分类训练

1、环境准备

2、数据准备

3、数据训练

六、接入 Swanlab 可视化训练结果

1、设置初始化配置参数

2、训练函数添加 Swanlab 跟踪

3、测试函数添加 Swanlab 跟踪

4、运行

七、使用 Gradio 进行功能演示

八、案例代码地址

参考文献：

一、简单介绍

Python是一种跨平台的计算机程序设计语言。是一种面向对象的动态类型语言，最初被设计用于编写自动化脚本(shell)，随着版本的不断更新和语言新功能的添加，越多被用于独立的、大型项目的开发。Python是一种解释型脚本语言，可以应用于以下领域： Web 和 Internet开发、科学计算和统计、人工智能、教育、桌面界面开发、软件开发、后端开发、网络爬虫。

Python 机器学习是利用 Python 编程语言中的各种工具和库来实现机器学习算法和技术的过程。Python 是一种功能强大且易于学习和使用的编程语言，因此成为了机器学习领域的首选语言之一。Python 提供了丰富的机器学习库，如Scikit-learn、TensorFlow、Keras、PyTorch等，这些库包含了许多常用的机器学习算法和深度学习框架，使得开发者能够快速实现、测试和部署各种机器学习模型。

通过 Python 进行机器学习，开发者可以利用其丰富的工具和库来处理数据、构建模型、评估模型性能，并将模型部署到实际应用中。Python 的易用性和庞大的社区支持使得机器学习在各个领域都得到了广泛的应用和发展。

二、PyTorch

PyTorch是一个开源的深度学习框架，以其出色的灵活性和易用性而著称。它是由Facebook AI Research及其他几个实验室的开发者共同开发的，将高效的GPU加速后端库Torch与直观的Python前端相结合，专注于快速原型设计、代码可读性，并支持广泛的深度学习模型。

使用PyTorch实现猫狗分类的原理

1、动态计算图

PyTorch使用动态计算图来定义和跟踪计算操作。与传统的静态计算图相比，动态计算图允许在执行过程中动态地构建计算图，使得可以使用常规的编程控制流语句（如if和while）来定义计算图的结构，实现更灵活和高效的模型构建和训练。

2、自动微分

PyTorch通过自动微分机制，实现了对计算图中各个操作的梯度自动计算。用户只需要定义计算图和前向传播过程，PyTorch会自动追踪计算图中的每个操作，并在需要时计算各个操作的梯度，大大简化了深度学习模型的训练过程。

3、张量计算

PyTorch的张量计算是其核心功能之一，提供了类似于NumPy的API，但可以在GPU上进行计算，从而加速大规模数值计算。

4、高效的并行计算

PyTorch支持高效的并行计算，可以利用多GPU进行训练，加速模型的训练过程。

5、构建模型的五要素

在PyTorch中实现猫狗分类，需要关注以下五个要素：

数据：包括数据读取、清洗、划分和预处理。
模型：包括构建模型模块、组织复杂网络、初始化网络参数和定义网络层。
损失函数：创建损失函数，设置超参数，并根据不同任务选择合适的损失函数。
优化器：根据梯度使用某种优化器更新参数，管理模型参数，调整学习率。
迭代训练：组织上述四个模块进行反复训练，观察训练效果，绘制Loss/Accuracy曲线或用TensorBoard进行可视化分析。

6、模型训练步骤

使用PyTorch实现猫狗分类的步骤通常包括：

数据预处理：包括数据增强，如随机裁剪、旋转、水平翻转等，以提高模型的泛化能力。
模型定义：定义CNN模型，可以是自定义的或基于预训练模型的结构。
损失函数与优化器：选择合适的损失函数（如交叉熵损失）和优化器（如Adam）。
训练循环：进行模型训练，包括前向传播、计算损失、反向传播和参数更新。
评估与测试：在独立的验证集和测试集上评估模型性能，使用准确率等指标。

通过上述原理和步骤，PyTorch提供了一个强大而灵活的平台，用于实现猫狗分类等深度学习任务。

三、CNN

1、神经网络

1.1 神经网络结构

可以通过下图进行理解神经网络的基本构成：

1.2 图片在计算机内的储存

图片在计算机储存由像素点矩阵组成，黑白图片的像素点是0-255或者0-1之间的数值，代表明暗程度；彩色图片是RGB图像，RGB表示红，绿，蓝三原色，计算机里所有的颜色都是三原色不同比例组成的，即三色通道

1.3 图像的传递

将二维图像经过flatten 展开成一维输入全连接网络中

1.4 训练数据

输入一组照片，通过全连接层的处理输出预测值和损失，损失越小越接近真实结果，因此需要找到最好的参数，即让所有的损失和最小，那么如何找到最好的参数呢？

现在选用的方法是梯度下降：

通过梯度下降不断迭代，调整初始参数，找到总损失比较小的最佳参数

2、卷积神经网络

2.1 图片的特质

2.1.1 图片的一些模式比整张图片小的多

比如说要识别猫，可以只通过猫的一部分特征去进行识别，即一个神经元不需要看到整个图像去发现模式，可以通过较少的参数连接到小区域

2.1.2 同样的模式可能出现在图像的不同区域

相同的猫耳检测器可以共享参数

2.1.3 对图像进行缩放不会改变图像中的物体

当图片很大时，图片的像素点也会很多，那么图片传入神经网络后连接数就会很多，参数就会多。缩放后可以使参数减少，简化问题

2.2 CNN模型

2.2.1 卷积层

卷积核在原始图片中起到探测模式的作用。可以发现卷积核的维度比原始图像要小，实现卷积的过程就是开始时，让卷积核从原始图像左上角对齐，对应每个小格子位置相乘，再将所有的结果相加，得到卷积结果矩阵的第一个值；再将卷积核向右移动，遍历原始图像，以此类推

不同的卷积核有不同的效果，而其中的值都是需要学习的参数

例：原始图片是8x8像素的，卷积核是3x3像素的，卷积结果是多少像素的？

答：6x6像素，8x8矩阵减去边缘一圈，即8-2=6

补充：

（1）边界处理

有两种边界处理方式，Full Padding和Same Padding

（2）Stride: 卷积核每次移动的步长

2.2.2 最大池化层

在每个小区域内最大值取出来组合，起到图像缩放的作用，减少参数

2.2.3 Flatten层

将二维图像经过flatten 展开成一维输入全连接层中

2.3 keras

Sequential 模型：非常简单，只支持单输入，单输出的模型，适用于70%的应用场景

函数式API：支持多输入，多输出模型，适用于95%的应用场景

建立一个全连接层：

import keras
from keras import layers#导入层结构model = keras.Sequential()  #建立序列模型
# 全连接层（本层神经元个数，激活函数，输入图片参数值数量）
model.add(layers.Dense(20, activation='relu', input_shape=(10,))) 
model.add(layers.Dense(20, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))# 训练模型
# x-样本数据即图片，y-图片标签，epochs=处理图片的次数，batch_size=一次性处理几张图片
model.fit(x, y, epochs=10, batch_size=32)

建立一个卷积层：

keras.layers.Conv2D(filters, kernel_size, strides=(1,1), padding='valid', data_format=None)
# filters: 输出空间的维度
# kernel_size: 1个整数或2个整数表示的元组，2D卷积窗口的宽度和高度
# strides: 2个整数表示的元组，卷积沿宽度和高度方向的步长
# padding: 边界处理的方法,"valid"或"same"

建立一个最大池化层：

keras.layers.MaxPooling2D(pool_size=(2,2), strides=None, padding='valid', data_format=None )
# pool_size: 沿（垂直，水平）方向缩小比例的因数，如果只有一个整数，则两个维度使用相同窗口长度
# strides: 2个整数表示的元组，步长值，None表示默认值pool_size
# padding: 边界处理的方法,"valid"或"same"

四、ResNet50

ResNet50是一种深度学习模型，属于残差网络（ResNet）家族，由微软研究院的Kaiming He等人于2015年提出。以下是关于ResNet50模型的详细说明：

1、ResNet50 的总体结构

ResNet50由多个卷积层、批量归一化层（Batch Normalization）、激活函数和残差块（Residual Block）组成，总共有50个卷积层。网络结构从输入到输出可以分为以下几个阶段：input->stage0->stage1->stage2->stage3->stage4->output。

2、残差块

ResNet50有两个基本的块，分别名为_Conv Block_和_Identity Block_。_Conv Block_用于改变网络的维度，而_Identity Block_用于加深网络

。残差块是ResNet的核心部分，它包含两个卷积层和一个快捷连接（Skip Connection），通过快捷连接，将输入直接加到卷积层的输出上，形成残差连接。这种结构使得网络在训练过程中能够更好地保留梯度信息，从而避免梯度消失的问题。

3、批量归一化层（Batch Normalization）

批量归一化层是一种常用于深度神经网络中的正则化技术，可以加速神经网络的训练过程，使得网络中的梯度在反向传播过程中更加稳定。

4、激活函数

ResNet50中通常使用ReLU作为激活函数。

5、参数数量

ResNet50的总参数数量约为25,636,712。

6、应用场景

ResNet50以其出色的图像识别能力而闻名，在图像分类、目标检测、图像分割等任务中取得了卓越的性能。通过加载预训练的ResNet50模型，并在特定数据集上进行微调，可以实现高效的模型训练和推理。

7、训练和优化

ResNet50的训练涉及数据集和预处理、超参数的设置和调整、训练过程的监控和可视化等方面。

8、评估和部署

模型评估指标和方法、模型部署的平台和工具也是ResNet50实际应用中需要考虑的重要方面。

五、Pytorch 实现猫狗分类训练

1、环境准备

案例环境：1） Windows 10；2）Python 3.11

构建虚拟环境，安装相关包，主要是：torch、torchvision、transforms

如果使用 cuda 进行训练，查看自己的 cuda 版本对应安装 torch 相关

案例中 cuda 版本为 12.3，所以对应安装 torch 如下命令：
pip install torch==2.0.0+cu118 torchvision==0.15.1+cu118 torchaudio==2.0.1+cu118 -f https://download.pytorch.org/whl/torch_stable.html

2、数据准备

2.1 数据源

可以去网上找例如：Mo-人工智能教学实训平台，在线学习Python、AI、大模型、AI写作绘画课程，零基础轻松入门

也可暂时使用该案例的数据

2.2 数据说明

说明：

在datasets目录下，train.csv和val.csv分别记录了训练集和测试集的图像相对路径（第一列是图像的相对路径，第二列是标签，0代表猫，1代表狗）：

2.3 DatasetLoader

创建DatasetLoader，主要功能是通过读取CSV文件来加载图像数据，并对其进行预处理，使其适合用于PyTorch模型的训练或测试。通过继承 torch.utils.data.Dataset 类，DatasetLoader 类实现了 __len__ 和 __getitem__ 方法，这两个方法是自定义数据集所必需的，分别用于获取数据集的大小和根据索引获取数据项。

import csv  # 导入csv模块，用于读取CSV文件
import os  # 导入os模块，用于处理文件路径
from torchvision import transforms  # 从torchvision导入transforms，用于图像预处理
from PIL import Image  # 从PIL库导入Image，用于图像读取和处理
from torch.utils.data import Dataset  # 从torch.utils.data导入Dataset，用于创建自定义数据集class DatasetLoader(Dataset):def __init__(self, csv_path):"""初始化函数，读取CSV文件并存储图像数据。参数:csv_path (str): CSV文件的路径，该文件包含图像路径和对应的标签。"""self.csv_file = csv_pathwith open(self.csv_file, 'r') as file:  # 打开CSV文件self.data = list(csv.reader(file))  # 读取CSV文件内容并存储为列表self.current_dir = os.path.dirname(os.path.abspath(__file__))  # 获取当前文件的目录路径def preprocess_image(self, image_path):"""预处理图像函数，对图像进行大小调整、转换为张量和归一化。参数:image_path (str): 图像的相对路径。返回:Tensor: 预处理后的图像张量。"""full_path = os.path.join(self.current_dir, 'datasets', image_path)  # 拼接完整的图像路径image = Image.open(full_path)  # 使用PIL库打开图像image_transform = transforms.Compose([  # 创建一个transforms.Compose对象，用于链式图像预处理transforms.Resize((256, 256)),  # 将图像大小调整为256x256transforms.ToTensor(),  # 将图像转换为PyTorch张量transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化处理])return image_transform(image)  # 返回预处理后的图像张量def __getitem__(self, index):"""获取数据项函数，根据索引获取图像和标签，并进行预处理。参数:index (int): 数据项的索引。返回:tuple: 包含预处理后的图像张量和标签的元组。"""image_path, label = self.data[index]  # 获取索引对应的图像路径和标签image = self.preprocess_image(image_path)  # 预处理图像return image, int(label)  # 返回图像张量和标签def __len__(self):"""获取数据集大小函数，返回数据集的总项数。返回:int: 数据集的总项数。"""return len(self.data)  # 返回数据集的总项数

DatasetLoader类由四个部分组成：

1）__init__：包含1个输入参数csv_path，在外部传入csv_path后，将读取后的数据存入self.data中。self.current_dir则是获取了当前代码所在目录的绝对路径，为后续读取图像做准备。

2）preprocess_image：此函数用于图像预处理。首先，它构造图像文件

的绝对路径，然后使用PIL库打开图像。接着，定义了一系列图像变换：调整图像大小至256x256、转换图像为张量、对图像进行标准化处理，最终，返回预处理后的图像。

3）__getitem__：当数据集类被循环调用时，__getitem__方法会返回指定索引index的数据，即图像和标签。首先，它根据索引从self.data中取出图像路径和标签。然后，调用prepogress_image方法来处理图像数据。最后，将处理后的图像数据和标签转换为整型后返回。

4）__len__：用于返回数据集的总图像数量。

3、数据训练

1、载入数据集

from torch.utils.data import DataLoader
from load_datasets import DatasetLoaderbatch_size = 8  # 设置批次大小TrainDataset = DatasetLoader("datasets/train.csv")  # 创建训练数据集对象
ValDataset = DatasetLoader("datasets/val.csv")  # 创建验证数据集对象
# 创建训练数据的DataLoader
TrainDataLoader = DataLoader(TrainDataset, batch_size=batch_size, shuffle=True) 
# 创建验证数据的DataLoader 
ValDataLoader = DataLoader(ValDataset, batch_size=1, shuffle=False)

这里传入那两个csv文件的路径实例化DatasetLoader类，然后用PyTorch的DataLoader做一层封装。 DataLoader可以再传入两个参数：

batch_size：定义了每个数据批次包含多少张图像。在深度学习中，我们通常不会一次性地处理所有数据，而是将数据划分为小批次。这有助于模型更快地学习，并且还可以节省内存。在这里我们定义batch_size = 8，即每个批次将包含8个图像。
shuffle：定义了是否在每个循环轮次（epoch）开始时随机打乱数据。这通常用于训练数据集以保证每个epoch的数据顺序不同，从而帮助模型更好地泛化。如果设置为True，那么在每个epoch开始时，数据将被打乱。在这里我们让训练时打乱，测试时不打乱。

2、载入ResNet50模型

模型选用经典的ResNet50，模型的具体原理本文不细说，重点放在工程实现上。

我们使用torchvision来创建1个resnet50模型，并载入预训练权重：

from torchvision.models import ResNet50_Weights# 加载预训练的ResNet50模型
model = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

因为猫狗分类是个2分类任务，而torchvision提供的resnet50默认是1000分类，所以我们需要把模型最后的全连接层的输出维度替换为2：

from torchvision.models import ResNet50_Weightsnum_classes=2 # 设置分类数# 加载预训练的ResNet50模型
model = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)# 将全连接层替换为2分类in_features = model.fc.in_featuresmodel.fc = torch.nn.Linear(in_features, num_classes)  # 修改全连接层为2个输出

3、设置cuda/mps/cpu

如果你的电脑是英伟达显卡，那么cuda可以极大加速你的训练；

如果你的电脑是Macbook Apple Sillicon（M系列芯片），那么mps

同样可以极大加速你的训练；

如果都不是，那就选用cpu：

# 设置device
try:use_mps = torch.backends.mps.is_available()  # 检查是否支持MPS
except AttributeError:use_mps = Falseif torch.cuda.is_available():  # 如果CUDA可用device = "cuda"
elif use_mps:  # 如果MPS可用device = "mps"
else:device = "cpu"  # 否则使用CPUprint("device is " + device)  # 打印使用的设备

将模型加载到对应的device中：

model.to(device)  # 将模型移动到指定设备

4、设置超参数、优化器、损失函数

设置训练轮次为20轮，学习率为1e-4，训练批次为8，分类数为2分类。

num_epochs = 20  # 设置总训练轮数
lr = 1e-4  # 设置学习率
batch_size = 8  # 设置批次大小
num_classes = 2  # 设置分类数

设置损失函数为交叉熵损失，优化器为Adam。

criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=lr)  # 定义优化器

5、训练函数

定义1个训练函数train

# 定义训练函数
def train(model, device, train_dataloader, optimizer, criterion, epoch, num_epochs):"""训练模型一个epoch。参数:model: 要训练的模型。device: 训练使用的设备（CPU或GPU）。train_dataloader: 训练数据的DataLoader。optimizer: 优化器。criterion: 损失函数。epoch: 当前epoch数。num_epochs: 总epoch数。"""model.train()  # 设置模型为训练模式for iter, (inputs, labels) in enumerate(train_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备optimizer.zero_grad()  # 清空梯度outputs = model(inputs)  # 前向传播loss = criterion(outputs, labels)  # 计算损失loss.backward()  # 反向传播optimizer.step()  # 更新参数# 打印训练信息print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, iter + 1, len(train_dataloader),loss.item()))

训练的逻辑：循环调用train_dataloader，每次取出1个batch_size的图像和标签，传入到resnet50模型中得到预测结果，将结果和标签传入损失函数中计算交叉熵损失，最后根据损失计算反向传播，Adam优化器执行模型参数更新，循环往复。

6、测试函数

定义1个测试函数test

# 定义测试函数
def test(model, device, test_dataloader, epoch):"""测试模型的性能。参数:model: 要测试的模型。device: 测试使用的设备（CPU或GPU）。test_dataloader: 测试数据的DataLoader。epoch: 当前epoch数。"""model.eval()  # 设置模型为评估模式correct = 0total = 0with torch.no_grad():  # 在这个上下文中，不计算梯度for iter, (inputs, labels) in enumerate(test_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备outputs = model(inputs)  # 前向传播_, predicted = torch.max(outputs.data, 1)  # 获取预测结果total += labels.size(0)  # 总样本数correct += (predicted == labels).sum().item()  # 计算正确预测的数量accuracy = correct / total * 100  # 计算准确率print('Accuracy: {:.2f}%'.format(accuracy))  # 打印准确率

测试的逻辑：循环调用test_dataloader，将测试集的图像传入到resnet50模型中得到预测结果，与标签进行对比，计算整体的准确率。

7、训练并测试，最后保存权重文件

# 开始训练
for epoch in range(1, num_epochs + 1):  # 遍历所有epochtrain(model, device, TrainDataLoader, optimizer, criterion, epoch, num_epochs)  # 训练一个epochif epoch % 4 == 0:  # 每4个epoch测试一次test(model, device, ValDataLoader, epoch)  # 测试模型性能# 保存权重文件
if not os.path.exists("checkpoint"):  # 如果checkpoint目录不存在，则创建os.makedirs("checkpoint")
torch.save(model.state_dict(), 'checkpoint/latest_checkpoint.pth')  # 保存模型权重
print("Training complete")  # 打印训练完成信息

8、运行脚本，训练结果

9、完整代码

import torch
import torchvision
from torchvision.models import ResNet50_Weights
from torch.utils.data import DataLoader
from load_datasets import DatasetLoader
import os# 定义训练函数
def train(model, device, train_dataloader, optimizer, criterion, epoch, num_epochs):"""训练模型一个epoch。参数:model: 要训练的模型。device: 训练使用的设备（CPU或GPU）。train_dataloader: 训练数据的DataLoader。optimizer: 优化器。criterion: 损失函数。epoch: 当前epoch数。num_epochs: 总epoch数。"""model.train()  # 设置模型为训练模式for iter, (inputs, labels) in enumerate(train_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备optimizer.zero_grad()  # 清空梯度outputs = model(inputs)  # 前向传播loss = criterion(outputs, labels)  # 计算损失loss.backward()  # 反向传播optimizer.step()  # 更新参数# 打印训练信息print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, iter + 1, len(train_dataloader),loss.item()))# 定义测试函数
def test(model, device, test_dataloader, epoch):"""测试模型的性能。参数:model: 要测试的模型。device: 测试使用的设备（CPU或GPU）。test_dataloader: 测试数据的DataLoader。epoch: 当前epoch数。"""model.eval()  # 设置模型为评估模式correct = 0total = 0with torch.no_grad():  # 在这个上下文中，不计算梯度for iter, (inputs, labels) in enumerate(test_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备outputs = model(inputs)  # 前向传播_, predicted = torch.max(outputs.data, 1)  # 获取预测结果total += labels.size(0)  # 总样本数correct += (predicted == labels).sum().item()  # 计算正确预测的数量accuracy = correct / total * 100  # 计算准确率print('Accuracy: {:.2f}%'.format(accuracy))  # 打印准确率if __name__ == "__main__":num_epochs = 20  # 设置总训练轮数lr = 1e-4  # 设置学习率batch_size = 8  # 设置批次大小num_classes = 2  # 设置分类数# 设置devicetry:use_mps = torch.backends.mps.is_available()  # 检查是否支持MPSexcept AttributeError:use_mps = Falseif torch.cuda.is_available():  # 如果CUDA可用device = "cuda"elif use_mps:  # 如果MPS可用device = "mps"else:device = "cpu"  # 否则使用CPUprint("device is " + device)  # 打印使用的设备TrainDataset = DatasetLoader("datasets/train.csv")  # 创建训练数据集对象ValDataset = DatasetLoader("datasets/val.csv")  # 创建验证数据集对象TrainDataLoader = DataLoader(TrainDataset, batch_size=batch_size, shuffle=True)  # 创建训练数据的DataLoaderValDataLoader = DataLoader(ValDataset, batch_size=1, shuffle=False)  # 创建验证数据的DataLoader# 载入ResNet50模型model = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)# 将全连接层替换为2分类in_features = model.fc.in_featuresmodel.fc = torch.nn.Linear(in_features, num_classes)  # 修改全连接层为2个输出model.to(device)  # 将模型移动到指定设备criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数optimizer = torch.optim.Adam(model.parameters(), lr=lr)  # 定义优化器# 开始训练for epoch in range(1, num_epochs + 1):  # 遍历所有epochtrain(model, device, TrainDataLoader, optimizer, criterion, epoch, num_epochs)  # 训练一个epochif epoch % 4 == 0:  # 每4个epoch测试一次test(model, device, ValDataLoader, epoch)  # 测试模型性能# 保存权重文件if not os.path.exists("checkpoint"):  # 如果checkpoint目录不存在，则创建os.makedirs("checkpoint")torch.save(model.state_dict(), 'checkpoint/latest_checkpoint.pth')  # 保存模型权重print("Training complete")  # 打印训练完成信息

六、接入 Swanlab 可视化训练结果

SwanLab是一个类似Tensorboard的开源训练图表可视化库，有着更轻量的体积与更友好的API。除了能记录指标，还能自动记录训练的logging、硬件环境、Python环境、训练时间等信息。

Swanlab 官网：SwanLab - AGI时代先进模型训练研发工具

Swanlab Github ：GitHub - SwanHubX/SwanLab: ⚡️SwanLab: your ML experiment notebook. 你的AI实验笔记本，日志记录与可视化AI训练全流程。

注意：记得 pip install swanlab 安装工具

1、设置初始化配置参数

import swanlab  # 导入SwanLab库，用于实验管理和可视化# 初始化SwanLab
swanlab.init(# 设置项目、实验名和实验介绍project="Cats_Dogs_Classification",experiment_name="ResNet50",description="用ResNet50训练猫狗分类任务",# 记录超参数config={"model": "resnet50","optim": "Adam","lr": lr,"batch_size": batch_size,"num_epochs": num_epochs,"num_class": num_classes,"device": device,},
)

2、训练函数添加 Swanlab 跟踪

# 定义训练函数
def train(model, device, train_dataloader, optimizer, criterion, epoch, num_epochs, TrainDataLoader):"""训练模型一个epoch。参数:model: 要训练的模型。device: 训练使用的设备（CPU或GPU）。train_dataloader: 训练数据的DataLoader。optimizer: 优化器。criterion: 损失函数。epoch: 当前epoch数。num_epochs: 总epoch数。TrainDataLoader: 训练数据的DataLoader，用于获取迭代次数。"""model.train()  # 设置模型为训练模式for iter, (inputs, labels) in enumerate(train_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备optimizer.zero_grad()  # 清空梯度outputs = model(inputs)  # 前向传播loss = criterion(outputs, labels)  # 计算损失loss.backward()  # 反向传播optimizer.step()  # 更新参数# 打印训练信息print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, iter + 1, len(TrainDataLoader),loss.item()))swanlab.log({"train_loss": loss.item()})  # 使用SwanLab记录训练损失

3、测试函数添加 Swanlab 跟踪

# 定义测试函数
def test(model, device, test_dataloader, epoch, class_name):"""测试模型的性能。参数:model: 要测试的模型。device: 测试使用的设备（CPU或GPU）。test_dataloader: 测试数据的DataLoader。epoch: 当前epoch数。class_name: 类别名称列表。"""model.eval()  # 设置模型为评估模式correct = 0total = 0with torch.no_grad():  # 在这个上下文中，不计算梯度images_list = []for iter, (inputs, labels) in enumerate(test_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备outputs = model(inputs)  # 前向传播_, predicted = torch.max(outputs.data, 1)  # 获取预测结果if iter < 30:images_list.append(swanlab.Image(inputs, caption=class_name[predicted.item()]))  # 使用SwanLab记录图像total += labels.size(0)  # 总样本数correct += (predicted == labels).sum().item()  # 计算正确预测的数量accuracy = correct / total * 100  # 计算准确率print('Accuracy: {:.2f}%'.format(accuracy))  # 打印准确率swanlab.log({"test_acc": accuracy})  # 使用SwanLab记录测试准确率swanlab.log({"Image": images_list})  # 使用SwanLab记录图像

4、运行

4.1 如果你第一次使用SwanLab，你需要先登录账号，在终端输入：

swanlab login

会让你填一个API Key，去SwanLab官网登录一下账号，在设置页面复制API Key，粘贴过来就可以：

4.2 运行脚本，运行结果

4.3 网页上的训练结果展示：

如图，看到train_loss和test_acc整体的变化曲线，以及我们测试集里的图像和它们对应的预测标签。

再切换到实验卡片，这里记录了实验的各种信息，包括超参数、最终的实验指标、实验状态、训练时长、Git仓库链接、主机名、操作系统、Python版本、硬件配置等等。

4.4 关键代码

import torch
import torchvision
from torchvision.models import ResNet50_Weights
import swanlab  # 导入SwanLab库，用于实验管理和可视化
from torch.utils.data import DataLoader
from load_datasets import DatasetLoader
import os# 定义训练函数
def train(model, device, train_dataloader, optimizer, criterion, epoch, num_epochs, TrainDataLoader):"""训练模型一个epoch。参数:model: 要训练的模型。device: 训练使用的设备（CPU或GPU）。train_dataloader: 训练数据的DataLoader。optimizer: 优化器。criterion: 损失函数。epoch: 当前epoch数。num_epochs: 总epoch数。TrainDataLoader: 训练数据的DataLoader，用于获取迭代次数。"""model.train()  # 设置模型为训练模式for iter, (inputs, labels) in enumerate(train_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备optimizer.zero_grad()  # 清空梯度outputs = model(inputs)  # 前向传播loss = criterion(outputs, labels)  # 计算损失loss.backward()  # 反向传播optimizer.step()  # 更新参数# 打印训练信息print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, iter + 1, len(TrainDataLoader),loss.item()))swanlab.log({"train_loss": loss.item()})  # 使用SwanLab记录训练损失# 定义测试函数
def test(model, device, test_dataloader, epoch, class_name):"""测试模型的性能。参数:model: 要测试的模型。device: 测试使用的设备（CPU或GPU）。test_dataloader: 测试数据的DataLoader。epoch: 当前epoch数。class_name: 类别名称列表。"""model.eval()  # 设置模型为评估模式correct = 0total = 0with torch.no_grad():  # 在这个上下文中，不计算梯度images_list = []for iter, (inputs, labels) in enumerate(test_dataloader):  # 遍历DataLoader中的批次inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到指定设备outputs = model(inputs)  # 前向传播_, predicted = torch.max(outputs.data, 1)  # 获取预测结果if iter < 30:images_list.append(swanlab.Image(inputs, caption=class_name[predicted.item()]))  # 使用SwanLab记录图像total += labels.size(0)  # 总样本数correct += (predicted == labels).sum().item()  # 计算正确预测的数量accuracy = correct / total * 100  # 计算准确率print('Accuracy: {:.2f}%'.format(accuracy))  # 打印准确率swanlab.log({"test_acc": accuracy})  # 使用SwanLab记录测试准确率swanlab.log({"Image": images_list})  # 使用SwanLab记录图像if __name__ == "__main__":num_epochs = 20  # 设置总训练轮数lr = 1e-4  # 设置学习率batch_size = 8  # 设置批次大小num_classes = 2  # 设置分类数# 设置devicetry:use_mps = torch.backends.mps.is_available()  # 检查是否支持MPSexcept AttributeError:use_mps = Falseif torch.cuda.is_available():  # 如果CUDA可用device = "cuda"elif use_mps:  # 如果MPS可用device = "mps"else:device = "cpu"  # 否则使用CPU# 初始化SwanLabswanlab.init(# 设置项目、实验名和实验介绍project="Cats_Dogs_Classification",experiment_name="ResNet50",description="用ResNet50训练猫狗分类任务",# 记录超参数config={"model": "resnet50","optim": "Adam","lr": lr,"batch_size": batch_size,"num_epochs": num_epochs,"num_class": num_classes,"device": device,},)TrainDataset = DatasetLoader("datasets/train.csv")  # 创建训练数据集对象ValDataset = DatasetLoader("datasets/val.csv")  # 创建验证数据集对象TrainDataLoader = DataLoader(TrainDataset, batch_size=batch_size, shuffle=True)  # 创建训练数据的DataLoaderValDataLoader = DataLoader(ValDataset, batch_size=1, shuffle=False)  # 创建验证数据的DataLoader# 载入ResNet50模型model = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)# 将全连接层替换为2分类in_features = model.fc.in_featuresmodel.fc = torch.nn.Linear(in_features, num_classes)  # 修改全连接层为2个输出model.to(torch.device(device))  # 将模型移动到指定设备criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数optimizer = torch.optim.Adam(model.parameters(), lr=lr)  # 定义优化器# 开始训练for epoch in range(1, num_epochs + 1):  # 遍历所有epochtrain(model, device, TrainDataLoader, optimizer, criterion, epoch, num_epochs, TrainDataLoader)  # 训练一个epochif epoch % 4 == 0:  # 每4个epoch测试一次test(model, device, ValDataLoader, epoch, ["cat", "dog"])  # 测试模型性能# 保存权重文件if not os.path.exists("checkpoint"):  # 如果checkpoint目录不存在，则创建os.makedirs("checkpoint")torch.save(model.state_dict(), 'checkpoint/latest_checkpoint.pth')  # 保存模型权重print("Training complete")  # 打印训练完成信息

七、使用 Gradio 进行功能演示

Gradio是一个开源的Python库，旨在帮助数据科学家、研究人员和从事机器学习领域的开发人员快速创建和共享用于机器学习模型的用户界面。

注意：记得使用 pip install gradio 进行工具安装

在这里我们使用Gradio来构建一个猫狗分类的Demo界面，编写app.py程序：

import gradio as gr
import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision# 加载与训练中使用的相同结构的模型
def load_model(checkpoint_path, num_classes):"""加载经过训练的模型。参数:checkpoint_path (str): 模型权重文件的路径。num_classes (int): 模型输出的类别数。返回:model: 加载了权重并设置为评估模式的模型。"""try:use_mps = torch.backends.mps.is_available()  # 检查是否支持MPS，MPS是苹果硬件上的Metal Performance Shadersexcept AttributeError:use_mps = False  # 如果不支持MPS，则设置为Falseif torch.cuda.is_available():  # 检查CUDA是否可用，即是否有NVIDIA GPUdevice = "cuda"  # 如果有NVIDIA GPU，则使用CUDAelif use_mps:  # 如果没有NVIDIA GPU但支持MPS，则使用MPSdevice = "mps"else:device = "cpu"  # 如果既没有NVIDIA GPU也不支持MPS，则使用CPUmodel = torchvision.models.resnet50(weights=None)  # 加载ResNet50模型，不加载预训练权重in_features = model.fc.in_features  # 获取全连接层的输入特征数model.fc = torch.nn.Linear(in_features, num_classes)  # 替换全连接层以匹配类别数model.load_state_dict(torch.load(checkpoint_path, map_location=device))  # 加载模型权重model.eval()  # 设置模型为评估模式return model# 加载图像并执行必要的转换的函数
def process_image(image, image_size):"""对图像进行预处理。参数:image: PIL图像对象。image_size (int): 图像的目标大小。返回:image: 预处理后的图像张量。"""# 定义与训练时相同的转换操作preprocessing = transforms.Compose([transforms.Resize((image_size, image_size)),  # 将图像大小调整为目标大小transforms.ToTensor(),  # 将图像转换为PyTorch张量transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化处理])image = preprocessing(image).unsqueeze(0)  # 增加批次维度return image# 预测图像类别并返回概率的函数
def predict(image):"""对输入图像进行分类并返回类别概率。参数:image: PIL图像对象。返回:class_probabilities: 类别概率字典。"""classes = {'0': 'cat', '1': 'dog'}  # 类别标签字典image = process_image(image, 256)  # 使用训练时的图像大小进行预处理with torch.no_grad():  # 不计算梯度，减少内存和计算资源消耗outputs = model(image)  # 模型前向传播probabilities = F.softmax(outputs, dim=1).squeeze()  # 应用softmax函数获取概率# 将类别标签映射到概率class_probabilities = {classes[str(i)]: float(prob) for i, prob in enumerate(probabilities)}return class_probabilities# 定义到您的模型权重的路径
checkpoint_path = 'checkpoint/latest_checkpoint.pth'
num_classes = 2
model = load_model(checkpoint_path, num_classes)  # 加载模型# 定义Gradio Interface
iface = gr.Interface(fn=predict,  # 绑定predict函数inputs=gr.Image(type="pil"),  # 输入为PIL图像outputs=gr.Label(num_top_classes=num_classes),  # 输出为类别标签title="Cat vs Dog Classifier",  # 界面标题
)if __name__ == "__main__":iface.launch()  # 启动Gradio界面

运行结果，拷贝网址，到浏览器上打开

在网页上打开，传图图片，效果如下

八、案例代码地址

https://download.csdn.net/download/u014361280/90071402

参考文献：

1、卷积神经网络-猫狗识别（附源码）_猫狗识别代码-CSDN博客

2、SwanHub - 创新的AI开源社区