狗的品种识别——基于PyTorch和ResNet

最新推荐文章于 2025-10-13 10:49:50 发布

原创最新推荐文章于 2025-10-13 10:49:50 发布 · 933 阅读 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

NSAcbba

关注

标签

#pytorch #人工智能 #python

分类人工智能

1. 项目介绍

在计算机视觉领域，图像分类是一个基础而重要的任务。本项目将使用PyTorch框架和ResNet模型来实现狗的品种识别系统。我们将使用斯坦福大学提供的"Stanford Dogs Dataset"数据集，该数据集包含120种不同品种的狗，共计20,580张图片。

本项目将完整展示从数据准备、模型构建、训练到评估的整个流程，并最终实现一个能够准确识别狗品种的深度学习模型。通过这个项目，读者可以学习到：

如何使用PyTorch处理图像数据
如何利用预训练的ResNet模型进行迁移学习
如何设计和优化深度学习模型
如何评估模型的性能

2. 环境准备
在开始项目前，我们需要准备以下环境和工具：

# 所需库的安装
!pip install torch torchvision
!pip install matplotlib numpy pandas
!pip install opencv-python
!pip install tqdm

# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
from PIL import Image
from tqdm import tqdm
import pandas as pd
import random

3. 数据集准备
3.1 数据集下载与探索
Stanford Dogs Dataset可以从官方网站下载，或者直接使用torchvision的API获取：

# 数据集路径
data_dir = './data/stanford_dogs'
if not os.path.exists(data_dir):
    os.makedirs(data_dir)
    
# 定义数据增强和标准化
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'valid': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# 加载数据集
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'valid']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=32,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'valid']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}
class_names = image_datasets['train'].classes

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

3.2 数据可视化

让我们看一下数据集中的一些样本：

def imshow(inp, title=None):
    """显示张量图像"""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # 暂停一下，以便更新绘图

# 获取一批训练数据
inputs, classes = next(iter(dataloaders['train']))

# 制作网格
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])
plt.show()

https://img-blog.csdnimg.cn/20210705163212345.png

4. 模型构建

4.1 ResNet模型简介

ResNet（Residual Neural Network）是由微软研究院的Kaiming He等人在2015年提出的深度卷积神经网络。它通过引入"残差学习"的概念，解决了深度神经网络中的梯度消失/爆炸问题，使得可以训练非常深的网络（如ResNet-152）。

本项目将使用预训练的ResNet-50模型进行迁移学习。

4.2 模型定义与修改

def initialize_model(num_classes, feature_extract=True, use_pretrained=True):
    """初始化这些变量，这些变量将在设置网络时设置。"""
    model_ft = models.resnet50(pretrained=use_pretrained)
    if feature_extract:
        # 冻结所有参数
        for param in model_ft.parameters():
            param.requires_grad = False
    
    # 修改最后的全连接层
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, num_classes)
    
    return model_ft

# 初始化模型
model_ft = initialize_model(len(class_names), feature_extract=True, use_pretrained=True)
model_ft = model_ft.to(device)

# 打印模型结构
print(model_ft)

4.3 训练参数设置

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()

# 只优化最后一层的参数
optimizer_ft = optim.SGD(model_ft.fc.parameters(), lr=0.001, momentum=0.9)

# 学习率调度器
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

5. 模型训练

5.1 训练函数定义

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # 每个epoch都有一个训练和验证阶段
        for phase in ['train', 'valid']:
            if phase == 'train':
                model.train()  # 设置模型为训练模式
            else:
                model.eval()   # 设置模型为评估模式

            running_loss = 0.0
            running_corrects = 0

            # 迭代数据
            for inputs, labels in tqdm(dataloaders[phase]):
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 梯度清零
                optimizer.zero_grad()

                # 前向传播
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 只在训练阶段进行反向传播和优化
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 统计
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            
            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            # 深度复制模型
            if phase == 'valid' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')

    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model

5.2 开始训练

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)

Epoch 0/24
----------
100%|██████████| 551/551 [02:15<00:00, 4.06it/s]
train Loss: 1.8214 Acc: 0.5823
100%|██████████| 138/138 [00:19<00:00, 7.06it/s]
valid Loss: 0.7561 Acc: 0.7872

Epoch 1/24
----------
100%|██████████| 551/551 [02:13<00:00, 4.14it/s]
train Loss: 0.8736 Acc: 0.7672
100%|██████████| 138/138 [00:19<00:00, 7.05it/s]
valid Loss: 0.5585 Acc: 0.8321
...
Training complete in 65m 23s
Best val Acc: 0.8724

6. 模型评估

6.1 测试集评估

def visualize_model(model, num_images=6):
    was_training = model.training
    model.eval()
    images_so_far = 0
    fig = plt.figure(figsize=(15, 10))

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['valid']):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//2, 2, images_so_far)
                ax.axis('off')
                ax.set_title(f'predicted: {class_names[preds[j]]}\nactual: {class_names[labels[j]]}')
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return
        model.train(mode=was_training)

visualize_model(model_ft)
plt.show()

6.2 混淆矩阵

from sklearn.metrics import confusion_matrix
import seaborn as sns

def plot_confusion_matrix(model, dataloader, class_names):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in tqdm(dataloader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(20, 20))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", 
                xticklabels=class_names, 
                yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.show()

plot_confusion_matrix(model_ft, dataloaders['valid'], class_names)

7. 模型优化与改进

7.1 微调所有层

之前我们只训练了最后一层，现在尝试微调所有层：

# 解冻所有层
for param in model_ft.parameters():
    param.requires_grad = True

# 重新定义优化器
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

# 再次训练
model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=10)

7.2 数据增强改进

增加更多的数据增强技术：

data_transforms_enhanced = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.RandomRotation(15),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'valid': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

7.3 使用不同的预训练模型

尝试使用ResNet-101或其他模型：

model_resnet101 = models.resnet101(pretrained=True)
num_ftrs = model_resnet101.fc.in_features
model_resnet101.fc = nn.Linear(num_ftrs, len(class_names))
model_resnet101 = model_resnet101.to(device)

8. 模型保存与部署

8.1 保存模型

torch.save(model_ft.state_dict(), 'dog_breed_classifier_resnet50.pth')

8.2 加载模型进行预测

def predict_breed(image_path, model, class_names, transform):
    image = Image.open(image_path)
    image = transform(image).unsqueeze(0).to(device)
    
    model.eval()
    with torch.no_grad():
        outputs = model(image)
        _, preds = torch.max(outputs, 1)
        prob = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
    
    return class_names[preds[0]], prob[preds[0]].item()

# 测试自定义图片
test_image = 'test_dog.jpg'
predicted_breed, confidence = predict_breed(
    test_image, 
    model_ft, 
    class_names, 
    data_transforms['valid']
)

print(f'Predicted breed: {predicted_breed} with confidence: {confidence:.2f}%')

# 显示测试图片
image = Image.open(test_image)
plt.imshow(image)
plt.title(f'Predicted: {predicted_breed} ({confidence:.2f}%)')
plt.axis('off')
plt.show()