深度学习入门：使用PyTorch和VGG16实现猫狗分类

一花·一树

已于 2024-02-23 13:12:25 修改

阅读量748

点赞数 13

文章标签：深度学习 pytorch 分类

于 2024-02-23 13:10:25 首次发布

本文链接：https://blog.csdn.net/qq_39045712/article/details/136253395

版权

摘要

深度学习已经成为计算机视觉领域解决分类问题的一种强大工具。本博客旨在为刚开始学习深度学习的朋友们提供一个简单的入门级示例：使用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 = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=True)

# 定义VGG16模型
class VGG16(nn.Module):
    def __init__(self, num_classes=6):
        super(VGG16, self).__init__()
        # 卷积层和池化层
        self.features = nn.Sequential(
            # conv1
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv4
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv5
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        # 全连接层
        self.classifier = nn.Sequential(
            nn.Linear(512 * 8 * 8, 4096),  # 512 * 8 * 8是根据网络结构和输入尺寸确定的
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平特征图
        x = self.classifier(x)
        return x


# 创建模型
model_net = VGG16(num_classes=2)

# 是否加载预训练模型，建议加载
pretrain = True 
if pretrain:
    # 加载预训练权重（确保替换为权重文件的实际路径）
    weights_path = './vgg16-Pretraining.pth'
    # 加载预训练模型的权重
    pretrained_dict = torch.load(weights_path)
    # 获取当前模型的state_dict
    model_dict = model_net.state_dict()
    # 过滤掉不匹配的权重
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
    # 更新当前模型的state_dict
    model_dict.update(pretrained_dict)
    # 加载我们真正需要的state_dict
    model_net.load_state_dict(model_dict)

model = model_net.to(device)

# 查看模型结构
from torchsummary import summary
summary(model, (3, 256, 256))

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

# 训练和验证模型
for epoch in range(EPOCH):  # 多次循环遍历数据集
    running_loss = 0.0
    model.train()  # 设置模型为训练模式
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs = model(inputs.float())
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 100 == 0:
            print(f"Epoch: {epoch + 1} / {EPOCH}    Step: {i} / {len(trainloader)}")

    # 验证阶段，验证运行时间较长
    correct = 0
    total = 0
    model.eval()  # 设置模型为评估模式
    with torch.no_grad():
        for data in valloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = model(images.float())
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}, Acc: {correct / total}')

# 保存训练好的模型权重
torch.save(model.state_dict(), "catdog_model.pth")
print('Finished Training')

训练过程展示：

在这里插入图片描述

模型推理

训练完成后，我们可以使用训练好的模型对新的图片进行分类预测。（cat_dog_classification_infer.py文件）

import torch  # 导入torch库，用于深度学习
import torchvision.transforms as transforms  # 导入torchvision库中的transforms模块，用于图像预处理
import torch.nn as nn  # 导入torch的神经网络模块
from PIL import Image  # 导入PIL库，用于图像处理
import matplotlib.pyplot as plt  # 导入matplotlib.pyplot模块，用于图形展示

# 设置设备为CUDA（如果可用），否则使用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device:", device)  # 打印当前使用的设备

# 定义数据预处理操作
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # 将图像大小调整为256x256
    transforms.ToTensor(),  # 将图像转换为Tensor
    transforms.Normalize((0.5,), (0.5,))  # 对图像进行标准化处理
])

# 定义VGG16模型类
class VGG16(nn.Module):
    def __init__(self, num_classes=6):  # 初始化函数，num_classes代表输出层类别数
        super(VGG16, self).__init__()  # 调用父类的初始化方法
        # 定义特征提取层（卷积层和池化层）
        self.features = nn.Sequential(
            # conv1
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv4
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # conv5
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        # 定义分类器（全连接层）
        self.classifier = nn.Sequential(
            nn.Linear(512 * 8 * 8, 4096),  # 全连接层
            nn.ReLU(inplace=True),  # 激活函数
            nn.Dropout(),  # Dropout层防止过拟合
            nn.Linear(4096, 4096),  # 全连接层
            nn.ReLU(inplace=True),  # 激活函数
            nn.Dropout(),  # Dropout层防止过拟合
            nn.Linear(4096, num_classes),  # 输出层
        )
    def forward(self, x):  # 定义前向传播过程
        x = self.features(x)  # 特征提取
        x = x.view(x.size(0), -1)  # 展平特征图
        x = self.classifier(x)  # 分类
        return x

# 创建模型实例并移至设定的设备
model_net = VGG16(num_classes=2)
model = model_net.to(device)

# 加载模型权重
model.load_state_dict(torch.load('catdog_model.pth', map_location=device))
model.to(device)  # 确保模型在正确的设备上
model.eval()  # 将模型设置为评估模式

# 加载并预处理图像
img = Image.open('./cat_dog_data/test/1.jpg')  # 加载预测的图像，更新为你的图片路径
img_transformed = transform(img)  # 应用预处理操作
img_batch = img_transformed.unsqueeze(0)  # 为图像添加批次维度以符合模型输入要求
# 进行预测
outputs = model(img_batch.to(device))
_, predicted = torch.max(outputs, 1)  # 获取预测结果中概率最高的类别

# 定义类别列表
classes = ['cat', 'dog']

# 展示图像及预测结果
plt.imshow(img)  # 展示图像
plt.title(f'Predicted: {classes[predicted[0]]}')  # 显示预测的类别
plt.show()  # 显示图形

结果展示：
预测结果展示

总结

这篇博客提供了使用PyTorch和VGG16进行猫狗分类的完整流程，包括数据处理、模型构建、训练、验证和推理。希望本教程能够为初学者提供一个良好的起点，鼓励他们在深度学习和人工智能的领域中继续探索和学习。

一花·一树

关注

13
点赞
踩
15

收藏

觉得还不错? 一键收藏
0
评论
深度学习入门：使用PyTorch和VGG16实现猫狗分类

这篇博客提供了使用PyTorch和VGG16进行猫狗分类的完整流程，包括数据处理、模型构建、训练、验证和推理。希望本教程能够为初学者提供一个良好的起点，鼓励他们在深度学习和人工智能的领域中继续探索和学习。
复制链接

扫一扫