一、工具准备

jupyter lab
python3.12
猫狗数据集准备

jupyter的结构

二、猫狗分类训练模型代码

1. 代码如下

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchsummary import summary
from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
import torch.nn.functional as F

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # 归一化
])

# 加载猫狗分类数据集
train_dataset = torchvision.datasets.ImageFolder(root='./data/catvsdog/train_dataset', transform=transform)
test_dataset = torchvision.datasets.ImageFolder(root='./data/catvsdog/test_dataset', transform=transform)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

# 定义一个简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 28 * 28, 500) # 表示该层将输入特征映射到 500 个特征
        self.fc2 = nn.Linear(500, 2)  # 2个类别（猫和狗）
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 64 * 28 * 28)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x


# 初始化模型、损失函数和优化器
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 打印模型结构
summary(model, (3, 224, 224))

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if (i+1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {running_loss/100:.3f}")
            running_loss = 0.0

print(f"GPU memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
# 保存模型的权重到文件
torch.save(model.state_dict(), './models/CatAndDog_Model.pth')

# 测试模型
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

# 打印分类报告
print(classification_report(all_labels, all_preds, target_names=['Cat', 'Dog']))

# 绘制混淆矩阵
cm = confusion_matrix(all_labels, all_preds)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['Cat', 'Dog'], yticklabels=['Cat', 'Dog'])
plt.xlabel('Predicted')
plt.ylabel('True')
plt.show()

2. 训练日志

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 16, 224, 224]             448
         MaxPool2d-2         [-1, 16, 112, 112]               0
            Conv2d-3         [-1, 32, 112, 112]           4,640
         MaxPool2d-4           [-1, 32, 56, 56]               0
            Conv2d-5           [-1, 64, 56, 56]          18,496
         MaxPool2d-6           [-1, 64, 28, 28]               0
            Linear-7                  [-1, 500]      25,088,500
           Dropout-8                  [-1, 500]               0
            Linear-9                    [-1, 2]           1,002
================================================================
Total params: 25,113,086
Trainable params: 25,113,086
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 13.41
Params size (MB): 95.80
Estimated Total Size (MB): 109.78
----------------------------------------------------------------
Epoch [1/10], Step [100/188], Loss: 0.750
Epoch [2/10], Step [100/188], Loss: 0.587
Epoch [3/10], Step [100/188], Loss: 0.530
Epoch [4/10], Step [100/188], Loss: 0.447
Epoch [5/10], Step [100/188], Loss: 0.412
Epoch [6/10], Step [100/188], Loss: 0.334
Epoch [7/10], Step [100/188], Loss: 0.266
Epoch [8/10], Step [100/188], Loss: 0.191
Epoch [9/10], Step [100/188], Loss: 0.143
Epoch [10/10], Step [100/188], Loss: 0.126
GPU memory allocated: 873.70 MB
              precision    recall  f1-score   support

         Cat       0.78      0.73      0.75      1000
         Dog       0.75      0.79      0.77      1000

    accuracy                           0.76      2000
   macro avg       0.76      0.76      0.76      2000
weighted avg       0.76      0.76      0.76      2000

三、验证训练出来的模型

1. 百度准备两张图片，狗和猫
   

2. 准备验证代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image

# 定义模型结构
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 28 * 28, 500)
        self.fc2 = nn.Linear(500, 2)  # 2个类别（猫和狗）
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 64 * 28 * 28)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 加载模型权重
model = SimpleCNN()
model.load_state_dict(torch.load('./models/CatAndDog_Model.pth', weights_only=False))
model.eval()  # 切换到评估模式

# 定义图片预处理流程
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # 归一化
])

# 加载图片
image_path = './data/images/cat2.jpg'  # 替换为你的图片路径
image = Image.open(image_path).convert('RGB')  # 确保图片是RGB格式
image = transform(image).unsqueeze(0)  # 添加批量维度

# 进行推理
with torch.no_grad():
    output = model(image)

# 获取预测结果
_, predicted = torch.max(output, 1)
predicted_label = predicted.item()

# 定义类别名称
class_names = ['Cat', 'Dog']

# 打印预测结果
print(f"预测结果: {class_names[predicted_label]}")

3. 执行结果

预测结果: Cat