摘要

深度学习已经成为计算机视觉领域解决分类问题的一种强大工具。本博客旨在为刚开始学习深度学习的朋友们提供一个简单的入门级示例：使用PyTorch框架和VGG16网络模型实现猫狗分类。我们将从数据准备、模型构建、训练过程、到最后的推理过程逐步讲解。

准备工作

在开始之前，请确保您的环境已安装以下依赖：

• torchsummary==1.5.1
• torch==1.13.0
• Pillow==8.2.0

您可以通过以下命令安装这些依赖：

pip install torchsummary==1.5.1 torch==1.13.0 Pillow==8.2.0

此外，您需要下载数据集、预训练的模型权重和代码文件。这些资源已经上传至百度网盘，链接如下：

·链接：https://pan.baidu.com/s/1lWgMRY40ksKHMkS8V3VICA
·提取码：w0r2

请下载并解压至您的工作目录。

数据集

数据集应包含两个文件夹：train 和 test，其中每个文件夹下分别存放猫和狗的图片。我们将使用train文件夹中的图片进行模型训练，并用test文件夹中的图片进行模型推理。

选择模型

我们将使用VGG16模型作为我们的基础架构。VGG16是一种流行的卷积神经网络架构，它在图像识别任务中表现出色。

构建和训练模型

首先，我们定义一个用于加载猫狗数据集的CatDogDataset类。这个类继承自torch.utils.data.Dataset，并实现了__len__和__getitem__方法，用于获取数据集大小和加载单个样本。在定义了数据加载和模型之后，我们可以开始训练模型了。训练过程包括多个epoch的迭代，每个epoch都会遍历整个训练集，并进行梯度下降以优化模型的权重。我们还会在每个epoch后在验证集上评估模型的性能，以监控模型是否过拟合。（cat_dog_classification_train.py文件）

import torch  # 导入PyTorch库
from PIL import Image  # 从PIL库导入Image模块，用于图像操作
from torch.utils.data import Dataset, DataLoader, random_split  # 导入数据集、数据加载器和随机分割功能
import os  # 导入os库，用于操作文件路径
import torch.optim as optim  # 导入优化器模块
import torch.nn as nn  # 导入神经网络模块

# 根据CUDA的可用性设置设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device:", device)  # 打印当前使用的设备

# 定义自己的数据集类
class CatDogDataset(Dataset):
    def __init__(self, directory, transform=None):
        self.directory = directory  # 数据集所在目录
        self.transform = transform  # 图像转换操作
        self.image_paths = []  # 存储图像路径
        self.labels = []  # 存储图像标签
        self.class_names = ['cat', 'dog']  # 类别名称
        # 遍历所有文件并收集路径和标签
        for label, class_name in enumerate(self.class_names):
            class_dir = os.path.join(self.directory, class_name)
            for img_file in os.listdir(class_dir):
                self.image_paths.append(os.path.join(class_dir, img_file))
                self.labels.append(label)

    def __len__(self):
        return len(self.image_paths)  # 返回数据集中图像的数量

    def __getitem__(self, index):
        # 根据索引获取图像路径和标签
        image_path = self.image_paths[index]
        label = self.labels[index]
        # 加载图像
        image = Image.open(image_path)
        # 如果指定了转换操作，则应用之
        if self.transform:
            image = self.transform(image)
        return image, label  # 返回图像和标签

import torchvision.transforms as transforms  # 导入图像转换模块

# 定义图像预处理步骤
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # 调整图像大小
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize((0.5,), (0.5,))  # 标准化
])
# 设置训练数据目录
train_dir = "./cat_dog_data/train"
# 创建数据集实例
dataset = CatDogDataset(train_dir, transform=transform)
# 分割数据集为训练集和验证集
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
print(f"train data: {
     
      
      len(train_dataset)}, val_data: {
     
      
      len(val_dataset)}")

# batch_size 每次训练加载图片数量
BATCH_SIZE = 128
# 训练迭代次数
EPOCH = 5

# 创建数据加载器
trainloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
valloader