使用ResNet18进行猫狗分类（原始数据处理+训练流程）

数据集处理脚本

原数据集下载地址：https://gitcode.com/open-source-toolkit/ef827/blob/main/cat_dog.zip
（其实用什么数据都行，这个数据集比较小）

import os
from shutil import copy
import random


def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)

#获取文件夹名，作为分类的类名
file_path = "cat_dog"
flower_class = [cla for cla in os.listdir(file_path)]

#创建 训练集train 文件夹，并由类名创建子目录
mkfile('data/train/')
for cla in flower_class:
    mkfile('data/train/' + cla)


mkfile('data/val/')
for cla in flower_class:
    mkfile('data/val/' + cla)

#划分比例：训练集:测试集 = 8:2
split_rate = 0.2

#遍历所有类别发全部图像并按比例划分
for cla in flower_class:
    cla_path = file_path + '/' + cla + '/'
    images = os.listdir(cla_path)
    num = len(images)
    eval_index = random.sample(images, k=int(num * split_rate))
    for index, image in enumerate(images):
        if image in eval_index:
            image_path = cla_path + image
            new_image_path = 'data/val/' + cla
            copy(image_path, new_image_path)
        else:
            image_path = cla_path + image
            new_image_path = 'data/train/' + cla
            copy(image_path, new_image_path)
        print('\r[{}] processed [{}/{}]'.format(cla, index+1, num))

print('DONE!')

解释

这个代码的主要功能是将位于指定文件夹中的猫狗图像分类到训练集和验证集，并根据一定的比例划分文件夹。这段代码如下是详细解析:

代码导入部分

import os
from shutil import copy
import random

• os：用于与操作系统交互，处理文件和目录的路径。
• shutil.copy：用于复制文件。
• random：用于生成随机数，这里用来随机选择某些图像。

mkfile 函数

def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)

• 定义一个函数 mkfile 用于创建一个文件夹。如果文件夹不存在，则创建它。

获取文件夹名作为类别名

file_path = "cat_dog"
flower_class = [cla for cla in os.listdir(file_path)]

• 定义了一个文件路径 file_path，其中存储有不同类别的猫狗图像（可能是不同的文件夹）。
• os.listdir(file_path) 获取指定路径下所有文件夹（类别）的名称，并储存在 flower_class 列表中。

创建训练集和验证集文件夹结构

#创建 训练集train 文件夹，并由类名创建子目录
mkfile('data/train/')
for cla in flower_class:
    mkfile('data/train/' + cla)

mkfile('data/val/')
for cla in flower_class:
    mkfile('data/val/' + cla)

• 创建 data/train/ 和 data/val/ 这两个主要文件夹，分别用于存放训练集和验证集。
• 随后，针对每个类别（类名），在这两个文件夹中各自创建一个以类名命名的子文件夹。

划分比例

#划分比例：训练集:测试集 = 8:2
split_rate = 0.2

• 设置训练集与验证集的划分比例为 8:2，即将 20% 的图像分配到验证集中。

遍历类别并划分图像

for cla in flower_class:
    cla_path = file_path + '/' + cla + '/'
    images = os.listdir(cla_path)
    num = len(images)
    eval_index = random.sample(images, k=int(num * split_rate))

• 针对每一个类别，构建完整的路径 cla_path 并获得该路径下的所有图像文件名。
• 计算图像的数量，并从中随机选择出 20% 的图像名称，储存在 eval_index 中，作为验证集的图像。

复制图像到相应的文件夹

for index, image in enumerate(images):
    if image in eval_index:
        image_path = cla_path + image
        new_image_path = 'data/val/' + cla
        copy(image_path, new_image_path)
    else:
        image_path = cla_path + image
        new_image_path = 'data/train/' + cla
        copy(image_path, new_image_path)
    print('\r[{}] processed [{}/{}]'.format(cla, index+1, num))

• 遍历每个图像，判断它是否在选中的验证集图像列表中。
• 如果是，则将这个图像复制到对应的验证集文件夹；如果不是，则复制到训练集文件夹。
• 使用 print 函数实时输出处理进度，显示当前类和处理的图像数量。

训练模型

1. 模型搭建

# %%
import os
import torch
import torch.nn as nn

# %%
from torchvision import models

class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18, self).__init__()
        # 加载预训练的 ResNet18 模型
        self.model = models.resnet18(pretrained=True)
        # 替换最后的全连接层，适配猫狗分类任务（输出2个类别）
        num_features = self.model.fc.in_features
        self.model.fc = nn.Linear(num_features, 2)

    def forward(self, x):
        return self.model(x)

说明

• 这里定义了一个自定义模型 ResNet18，继承自 nn.Module。
• 加载了 torchvision 提供的预训练 ResNet18 模型，并将最后的全连接层替换为适配当前分类任务的层（输出类别数为 2，分别表示猫和狗）。（当然）
• 在 forward 方法中，直接调用修改后的 ResNet 模型完成前向传播。

当然也可以根据ResNet的形状来自己一层层搭建,就像这样。（不过本文使用的是torchvision的ResNet18 模型）

import torch.nn.functional as F  

class BasicBlock(nn.Module):  
    """一个基本的残差块，使用两个3x3卷积层。"""  
    
    expansion = 1  # 用于控制输出通道数的扩展  

    def __init__(self, in_channels, out_channels, stride=1, downsample=None):  
        super(BasicBlock, self).__init__()  
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)  
        self.bn1 = nn.BatchNorm2d(out_channels)  
        self.relu = nn.ReLU(inplace=True)  
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)  
        self.bn2 = nn.BatchNorm2d(out_channels)  
        self.downsample = downsample  

    def forward(self, x):  
        identity = x  # 保存输入，以便计算残差  

        out = self.conv1(x)  
        out = self.bn1(out)  
        out = self.relu(out)  

        out = self.conv2(out)  
        out = self.bn2(out)  

        if self.downsample is not None:  
            identity = self.downsample(x)  # 如果需要对输入进行下采样  

        out += identity  # 残差连接  
        out = self.relu(out)  

        return out  


class ResNet(nn.Module):  
    """完整的 ResNet 架构。"""  

    def __init__(self, block, layers, num_classes=2):  
        super(ResNet, self).__init__()  
        self.in_channels = 64  
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)  
        self.bn1 = nn.BatchNorm2d(64)  
        self.relu = nn.ReLU(inplace=True)  
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  

        self.layer1 = self._make_layer(block, 64, layers[0])  # 第一层  
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)  # 第二层  
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)  # 第三层  
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)  # 第四层  
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # 自适应平均池化  
        self.fc = nn.Linear(512 * block.expansion, num_classes)  # 分类层  

    def _make_layer(self, block, out_channels, blocks, stride=1):  
        downsample = None  
        if stride != 1 or self.in_channels != out_channels * block.expansion:  
            downsample = nn.Sequential(  
                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),  
                nn.BatchNorm2d(out_channels * block.expansion),  
            )  

        layers = []  
        layers.append(block(self.in_channels, out_channels, stride, downsample))  
        self.in_channels = out_channels * block.expansion  
        for _ in range(1, blocks):  
            layers.append(block(self.in_channels, out_channels))  

        return nn.Sequential(*layers)  

    def forward(self, x):  
        x = self.conv1(x)  
        x = self.bn1(x)  
        x = self.relu(x)  
        x = self.maxpool(x)  

        x = self.layer1(x)  
        x = self.layer2(x)  
        x = self.layer3(x)  
        x = self.layer4(x)  

        x = self.avgpool(x)  
        x = torch.flatten(x, 1)  
        x = self.fc(x)  

        return x  

x = torch.randn(2, 3, 224, 224)  # 创建一个随机输入张量 (batch_size=2, channe
```ls=3, height=224, width=224)
model = ResNet18()  # 实例化模型
y = model(x)  # 前向传播
y.shape

说明

• 测试模型的初始化和前向传播是否工作正常。
• 输入张量大小为 (2, 3, 224, 224)，代表 2 张 224×224 的 RGB 图片。
• 输出大小为 (2, 2)，代表 2 个样本对应的 2 个类别的预测结果。

2. 数据准备

# %%
from torch.optim import lr_scheduler
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader

root_train = "data/train"  # 训练集路径
root_test = "data/val"  # 验证集路径

说明

• 使用 PyTorch 提供的 datasets.ImageFolder 加载本地数据集。
• 数据文件夹结构要求：root/class_name/image.jpg，不同类别的图片放置在独立文件夹中。

# 将图像的像素值归一化到 [-1, 1] 范围
# 使用 ResNet 预训练模型的标准化参数
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

# 定义训练数据增强
train_transform = transforms.Compose([
    transforms.Resize((224,224)),  # 调整图像大小
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.RandomRotation(15),  # 随机旋转
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 调整图像颜色
    transforms.RandomAffine(5, translate=(0.1, 0.1)),  # 随机仿射变换
    transforms.ToTensor(),  # 转为张量
    normalize])

# 定义验证数据预处理
val_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    normalize])

说明

• 训练集应用了多种数据增强方法，旨在提高模型的泛化能力。（说人话就是让模型能够认识更抽象的猫狗照片）
• 验证集只进行必要的预处理（如调整大小和归一化），以保证评估的公平性。

# 创建数据加载器
train_dataset = datasets.ImageFolder(root_train, train_transform)
val_dataset = datasets.ImageFolder(root_test, val_transform)

batch_size = 64
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)

说明

• ImageFolder 根据文件夹结构自动加载数据，并为每个类别分配标签。这里就导入了之前用脚本划分好的数据集。
• 使用 DataLoader 打包数据集，支持按批次加载和随机打乱。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = ResNet18().to(device)  # 模型加载到 GPU 或 CPU

# 定义损失函数和优化器
loss_fn = nn.CrossEntropyLoss()  # 多分类交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 学习率调度器：每隔 10 个 epoch 将学习率衰减到原来的 0.5
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

# 定义训练 epoch 数
epoch = 20

说明

• 使用 Adam 优化器，初始学习率为 0.0001。（以下是AI的解释）

• Adam（Adaptive Moment Estimation）是一种自适应学习率优化算法，它结合了动量（momentum）和 RMSProp 的优点。Adam 能够根据每个参数的稀疏性自适应地调整学习率，这使得它在许多任务中表现良好。

• 学习率设置为 1e-4：这个值是一个常用的初始学习率，尤其是在使用 Adam 优化器时。通常，学习率的选择是经验性的，1e-4 是许多研究和实践中发现的一个不错的起始值。你可以根据模型的训练情况（如损失是否收敛、准确率是否提升等）来调整学习率。

• 引入学习率调度器，避免过早收敛和过拟合。

• StepLR 是一种常见的调度策略，它在每隔一定的 epoch 后（这里是 10 个 epoch）将学习率乘以一个因子（这里是 0.5）。

依据：

1. 减小学习率：在训练的初期，较大的学习率有助于快速找到较优的解，但随着训练的深入，较小的学习率可以帮助模型更精细地调整参数，避免在局部最优解附近震荡。
2. 经验法则：很多情况下，随着训练轮次的增加，模型的性能提升会逐渐减缓，因此在达到一定的 epoch 后，降低学习率是一个常见的做法。通过减小学习率，可以使得模型在接近最优解时更稳定地收敛。
3. 调整学习率的策略：选择每 10 个 epoch 调整一次学习率和将其减小一半的策略是基于经验的，具体数值可以根据实际情况进行调整。一般来说，初始学习率和调度策略的选择可以通过实验来优化。

3. 模型训练和验证

训练函数

def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    for batch_idx, (X, targets) in enumerate(dataloader):
        images, targets = X.to(device), targets.to(device)
        outputs = model(images)  # 前向传播
        cur_loss = loss_fn(outputs, targets)  # 计算损失
        _, pred = outputs.max(dim=1)  # 获取预测值
        cur_acc = torch.sum(pred == targets) / outputs.shape[0]  # 计算准确率
        
        # 反向传播与优化
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()

        loss += cur_loss.item()
        current += cur_acc.item()
        n += 1
        
    train_acc = current / n
    train_loss = loss / n
    print("train loss: %.4f, train acc: %.4f" % (train_loss, train_acc))
    return train_loss, train_acc

验证函数

def val(dataloader, model, loss_fn, optimizer):
    model.eval()  # 设置模型为验证模式
    loss, current, n = 0.0, 0.0, 0
    with torch.no_grad():
        for batch_idx, (X, targets) in enumerate(dataloader):
            images, targets = X.to(device), targets.to(device)
            outputs = model(images)
            cur_loss = loss_fn(outputs, targets)
            _, pred = outputs.max(dim=1)
            cur_acc = torch.sum(pred == targets) / outputs.shape[0]
            
            loss += cur_loss.item()
            current += cur_acc.item()
            n += 1

    val_acc = current / n
    val_loss = loss / n
    print("val loss: %.4f, val acc: %.4f" % (val_loss, val_acc))
    return val_loss, val_acc

4. 训练过程

#这四个列表用来存储每一次的损失值和精确值，用于画图使用
# 也可以用d2l里的Animator进行实时画图
loss_train_list = []
acc_train_list = []
loss_val_list = []
acc_val_list = []

# 定义模型保存路径
folder = 'save_model'  
if not os.path.exists(folder):  
    os.makedirs(folder)  

max_acc = 0  # 初始化最大验证准确率
for t in range(epoch):  
    lr_scheduler.step()  
    print('Epoch %d/%d------' % (t + 1, epoch))  
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer)  
    val_loss, val_acc = val(val_dataloader, model, loss_fn, optimizer)  

    loss_train_list.append(train_loss)
    acc_train_list.append(train_acc)
    loss_val_list.append(val_loss)
    acc_val_list.append(val_acc)
    
    if val_acc > max_acc:  
        max_acc = val_acc  
        torch.save(model, os.path.join(folder, 'resnet18_best.pt'))  # 保存最佳模型
        print(f'Saving model, 第{t + 1}轮...')  

torch.save(model, os.path.join(folder, 'resnet18_last.pt'))  # 保存最后一次模型
print('最高精确值:', max_acc)

绘制训练曲线

import matplotlib.pyplot as plt

def plot_training_curves(loss_train_list, acc_train_list, loss_val_list, acc_val_list):
    """
    绘制训练过程中的损失值曲线和精确度曲线
    """
    epochs = range(1, len(loss_train_list) + 1)  # x 轴：训练轮次

    # 绘制损失值曲线
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)  # 第一张子图
    plt.plot(epochs, loss_train_list, label="Training Loss", marker='o', color='b')
    plt.plot(epochs, loss_val_list, label="Validation Loss", marker='o', color='r')
    plt.title("Loss Curve")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)

    # 绘制精确度曲线
    plt.subplot(1, 2, 2)  # 第二张子图
    plt.plot(epochs, acc_train_list, label="Training Accuracy", marker='o', color='g')
    plt.plot(epochs, acc_val_list, label="Validation Accuracy", marker='o', color='orange')
    plt.title("Accuracy Curve")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()
    plt.grid(True)

    # 显示图形
    plt.tight_layout()
    plt.show()

plot_training_curves(loss_train_list, acc_train_list, loss_val_list, acc_val_list)

可以看到画出来的曲线比较合理。

测试模型

import random

# 定义一个函数，绘制图像和标题
def show_images_with_predictions(imgs, preds, labels, num_rows, num_cols, classes, scale=1.5):
    """绘制图像列表，并显示预测结果与真实标签"""
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()  # 将 axes 从二维数组转换为一维数组
    for i, (ax, img, pred, label) in enumerate(zip(axes, imgs, preds, labels)):
        # 反归一化后转换为 numpy
        img = img.permute(1, 2, 0).numpy()  # (C, H, W) -> (H, W, C)
        ax.imshow(img)
        ax.axes.get_xaxis().set_visible(False)  # 隐藏 X 轴
        ax.axes.get_yaxis().set_visible(False)  # 隐藏 Y 轴
        ax.set_title(f"Truth: {classes[label]} \nPred: {classes[pred]}")
    plt.tight_layout()
    plt.show()

# 加载已训练的模型
model = torch.load('/save_model/resnet18_best.pt')
model = model.to(device)
classes = ['cat', 'dog']

model.eval()
imgs, preds, labels = [], [], []

# 批量处理数据
with torch.no_grad():
    for _ in range(10):  # 获取前 10 张图片
        i=random.randint(0, len(val_dataset)-1)
        img, label = val_dataset[i][0], val_dataset[i][1]
        img_with_batch = img.unsqueeze(0).to(device)  # 添加 batch 维度
        output = model(img_with_batch)  # 模型推理
        pred = torch.max(output, 1)[1].item()  # 获取预测类别索引

        imgs.append(img)  # 原始图像（张量）
        preds.append(pred)  # 预测结果
        labels.append(label)  # 真实标签

# 显示图片及预测结果
num_rows = 2
num_cols = 5
show_images_with_predictions(imgs, preds, labels, num_rows, num_cols, classes)

加载之前保存的训练好的模型.pt文件，然后进行测试。这里是随机从测试数据集中选取了10张图片进行检验。

完整训练代码

# %%
import os
import torch
import torch.nn as nn


# %% [markdown]
# ## 模型搭建

# %%
from torchvision import models

class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18, self).__init__()
        # 加载预训练的 ResNet18 模型
        self.model = models.resnet18(pretrained=True)
        # 替换最后的全连接层，适配猫狗分类任务（输出2个类别）
        num_features = self.model.fc.in_features
        self.model.fc = nn.Linear(num_features, 2)

    def forward(self, x):
        return self.model(x)


# %%
x=torch.randn(2,3,224,224)
model = ResNet18()
y = model(x)
y.shape

# %% [markdown]
# ## 训练

# %%
from torch.optim import lr_scheduler
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader

root_train = "data/train"
root_test = "data/val"

#将图像的像素值归一化到[-1,1]之间
#由于使用的是 torchvision.models 的预训练模型，所以用 ResNet 训练时的标准化参数
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])


train_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15), # 随机旋转
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1), # 调整图像颜色
    transforms.RandomAffine(5, translate=(0.1, 0.1)),  # 随机仿射变换
    transforms.ToTensor(),
    normalize])

val_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    normalize])


train_dataset = datasets.ImageFolder(root_train, train_transform)
val_dataset = datasets.ImageFolder(root_test, val_transform)

batch_size=64
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = ResNet18().to(device)

#定义损失函数
loss_fn = nn.CrossEntropyLoss()

#定义优化器（Adam）
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)


#调整学习率，每隔10epoch，变为原来的0.5
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

#训练批次
epoch=20


# %% [markdown]
# 定义训练函数

# %%
def train(dataloader, model, loss_fn, optimizer):
    loss,current,n = 0.0, 0.0, 0
    for batch_idx, (X, targets) in enumerate(dataloader):
        images, targets = X.to(device), targets.to(device)
        outputs = model(images)
        cur_loss = loss_fn(outputs, targets)
        _, pred = outputs.max(dim=1)
        cur_acc = torch.sum(pred == targets)/ outputs.shape[0]
        
        #反向传播
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()
        loss += cur_loss.item()
        current += cur_acc.item()
        n = n + 1
        
        
    train_acc = current / n
    train_loss = loss / n

    
    print("train loss: %.4f, train acc: %.4f" % (train_loss, train_acc))
    return train_loss, train_acc


# %% [markdown]
# 定义验证函数

# %%
def val(dataloader, model, loss_fn, optimizer):
    model.eval()
    loss,current,n = 0.0, 0.0, 0
    with torch.no_grad():
        for batch_idx, (X, targets) in enumerate(dataloader):
            images, targets = X.to(device), targets.to(device)
            outputs = model(images)
            cur_loss = loss_fn(outputs, targets)
            _, pred = outputs.max(dim=1)
            cur_acc = torch.sum(pred == targets)/ outputs.shape[0]
            loss += cur_loss.item()
            current += cur_acc.item()
            
            n = n + 1
        
    val_acc = current / n
    val_loss = loss / n
    print("val loss: %.4f, val acc: %.4f" % (val_loss, val_acc))
    return val_loss, val_acc

# %% [markdown]
# 查看模型在训练中的梯度和权重分布

# %% [markdown]
# 开始训练

# %%
loss_train_list=[]
acc_train_list=[]
loss_val_list=[]
acc_val_list=[]


# 定义模型保存的文件夹  
folder = 'save_model'  
if not os.path.exists(folder):  
    os.makedirs(folder)  

max_acc = 0  
for t in range(epoch):  
    lr_scheduler.step()  
    print('Epoch %d/%d------' % (t + 1, epoch))  
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer)  
    val_loss, val_acc = val(val_dataloader, model, loss_fn, optimizer)  

    
    loss_train_list.append(train_loss)
    acc_train_list.append(train_acc)
    loss_val_list.append(val_loss)
    acc_val_list.append(val_acc)
    
    if val_acc > max_acc:  
        max_acc = val_acc  
        torch.save(model, os.path.join(folder, 'resnet18_best.pt'))  
        print(f'Saving model, 第{t + 1}轮...')  
        
# 在最后一次训练结束后保存模型  
torch.save(model, os.path.join(folder, 'resnet18_last.pt'))  

print('最高精确值:',max_acc)
print("Done!")  

# %% [markdown]
# #绘制训练曲线

# %%
import matplotlib.pyplot as plt

def plot_training_curves(loss_train_list, acc_train_list, loss_val_list, acc_val_list):
    """
    绘制训练过程中的损失值曲线和精确度曲线
    """
    epochs = range(1, len(loss_train_list) + 1)  # x 轴：训练轮次

    # 绘制损失值曲线
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)  # 第一张子图
    plt.plot(epochs, loss_train_list, label="Training Loss", marker='o', color='b')
    plt.plot(epochs, loss_val_list, label="Validation Loss", marker='o', color='r')
    plt.title("Loss Curve")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)

    # 绘制精确度曲线
    plt.subplot(1, 2, 2)  # 第二张子图
    plt.plot(epochs, acc_train_list, label="Training Accuracy", marker='o', color='g')
    plt.plot(epochs, acc_val_list, label="Validation Accuracy", marker='o', color='orange')
    plt.title("Accuracy Curve")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()
    plt.grid(True)

    # 显示图形
    plt.tight_layout()
    plt.show()


# %%
plot_training_curves(loss_train_list, acc_train_list, loss_val_list, acc_val_list)


# %% [markdown]
# 测试模型

# %%
import random

# 定义一个函数，绘制图像和标题
def show_images_with_predictions(imgs, preds, labels, num_rows, num_cols, classes, scale=1.5):
    """绘制图像列表，并显示预测结果与真实标签"""
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()  # 将 axes 从二维数组转换为一维数组
    for i, (ax, img, pred, label) in enumerate(zip(axes, imgs, preds, labels)):
        # 反归一化后转换为 numpy
        img = img.permute(1, 2, 0).numpy()  # (C, H, W) -> (H, W, C)
        ax.imshow(img)
        ax.axes.get_xaxis().set_visible(False)  # 隐藏 X 轴
        ax.axes.get_yaxis().set_visible(False)  # 隐藏 Y 轴
        ax.set_title(f"Truth: {classes[label]} \nPred: {classes[pred]}")
    plt.tight_layout()
    plt.show()

# 加载已训练的模型
model = torch.load('save_model/resnet18_best.pt')
model = model.to(device)
classes = ['cat', 'dog']

model.eval()
imgs, preds, labels = [], [], []

# 批量处理数据
with torch.no_grad():
    for _ in range(10):  # 获取前 10 张图片
        i=random.randint(0, len(val_dataset)-1)
        img, label = val_dataset[i][0], val_dataset[i][1]
        img_with_batch = img.unsqueeze(0).to(device)  # 添加 batch 维度
        output = model(img_with_batch)  # 模型推理
        pred = torch.max(output, 1)[1].item()  # 获取预测类别索引

        imgs.append(img)  # 原始图像（张量）
        preds.append(pred)  # 预测结果
        labels.append(label)  # 真实标签

# 显示图片及预测结果
num_rows = 2
num_cols = 5
show_images_with_predictions(imgs, preds, labels, num_rows, num_cols, classes)