任务描述
猫狗分类任务（Dogs vs Cats）是Kaggle平台在2013年举办的一个经典计算机视觉竞赛。官方给出的Kaggle Dogs vs Cats 数据集中包括由12500张猫咪图片和12500张狗狗图片组成的训练集，12500张未标记照片组成的测试集。选手需要在规定时间内搭建模型，并在测试集中取得尽可能高的准确率。
数据集中的图片长这样👆，是不是还蛮贴近生活呢~
本文将设计一个简单的三层卷积神经网络，完成分类任务。基于Pyqt5设计交互界面，在测试集中检验分类结果。工作量不大非常适合新手作为入门项目~最终功能实现如下图所示。
代码介绍
我觉得新手最有成就感的事情莫过于跑通一个功能了，我们抛开原理不谈，先完整运行程序。代码我是用pycharm写的，新建一个classfy_program项目，在classfy_program项目中创建两个py文件，分别用于训练和测试数据集。此外还创建data文件夹，用于存放训练和测试用的图片。
cat_dog_test.py代码：
import sys
import torch
import torch.nn as nn
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QVBoxLayout, QFileDialog
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt
from torchvision import transforms
from PIL import Image
# ---------------- 定义模型结构 ----------------
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(128 * 28 * 28, 512)
self.fc2 = nn.Linear(512, 2) # 猫 vs 狗
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = self.pool(torch.relu(self.conv3(x)))
x = x.view(-1, 128 * 28 * 28)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# ---------------- 图像预处理 ----------------
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
# 类别名称（根据你的训练数据文件夹顺序）
class_names = ['猫', '狗']
# ---------------- 主窗口 ----------------
class CatDogClassifier(QWidget):
def __init__(self):
super().__init__()
self.setWindowTitle("猫狗分类器")
self.setGeometry(100, 100, 400, 500)
# 加载模型
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = CNN().to(self.device)
self.model.load_state_dict(torch.load("model_checkpoint.pth", map_location=self.device))
self.model.eval()
# 创建界面组件
self.image_label = QLabel("请点击按钮加载图片", self)
self.image_label.setAlignment(Qt.AlignCenter)
self.image_label.setFixedSize(300, 300)
self.result_label = QLabel("", self)
self.result_label.setAlignment(Qt.AlignCenter)
self.load_button = QPushButton("加载图片", self)
self.load_button.clicked.connect(self.load_image)
# 布局
layout = QVBoxLayout()
layout.addWidget(self.image_label)
layout.addWidget(self.result_label)
layout.addWidget(self.load_button)
self.setLayout(layout)
def load_image(self):
file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Image Files (*.png *.jpg *.jpeg)")
if file_path:
# 显示图片
pixmap = QPixmap(file_path)
pixmap = pixmap.scaled(300, 300, Qt.KeepAspectRatio)
self.image_label.setPixmap(pixmap)
# 进行预测
image = Image.open(file_path).convert("RGB")
input_tensor = transform(image).unsqueeze(0).to(self.device)
with torch.no_grad():
output = self.model(input_tensor)
_, predicted = torch.max(output, 1)
result = class_names[predicted.item()]
self.result_label.setText(f"预测结果：{result}")
# ---------------- 运行主程序 ----------------
if __name__ == "__main__":
app = QApplication(sys.argv)
window = CatDogClassifier()
window.show()
sys.exit(app.exec_())
cat_dog_train.py代码：
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 数据预处理与加载
transform = transforms.Compose([
transforms.Resize((224,224)), # 调整图片尺寸为224x224
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # 适用于自定义模型的标准化
])
# 设置训练集和验证集
train_dataset = datasets.ImageFolder(root='data/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# 自定义卷积神经网络模型
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
# 第一层卷积
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) # 输入通道为3（RGB），输出通道为32
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
# 池化层
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
# 全连接层
self.fc1 = nn.Linear(128 * 28 * 28, 512) # 调整为经过池化后的维度
self.fc2 = nn.Linear(512, 2) # 输出为2（猫 vs 狗）
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x))) # 第一层卷积 + ReLU + 最大池化
x = self.pool(torch.relu(self.conv2(x))) # 第二层卷积 + ReLU + 最大池化
x = self.pool(torch.relu(self.conv3(x))) # 第三层卷积 + ReLU + 最大池化
x = x.view(-1, 128 * 28 * 28) # 展平
x = torch.relu(self.fc1(x)) # 全连接层 + ReLU
x = self.fc2(x) # 输出层
return x
# 实例化模型
model = CNN()
# 使用GPU训练（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device) # 将模型移动到GPU（如果可用）
# 损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)
# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0
# 训练阶段
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device) # 将数据移动到GPU（如果可用）
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
# 输出每轮训练损失与准确率
print(
f"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader):.4f}, Accuracy: {100 * correct / total:.2f}%")
# 保存模型检查点
torch.save(model.state_dict(), 'model_checkpoint.pth')
print("Training complete!")
data数据集的下载链接在这里了：
Download Kaggle Cats and Dogs Dataset from Official Microsoft Download Center
希望看到这里的大家能够点一个小小的赞❤❤，后续会持续更新更多内容~~
