分享一套锋哥原创的基于PyTorch的猫狗图像识别系统(深度学习+PyQt6+ResNet18+ImageNet+迁移学习)

大家好，我是Java1234_小锋老师，分享一套锋哥原创的基于PyTorch的猫狗图像识别系统(深度学习+PyQt6+ResNet18+ImageNet+迁移学习)

项目介绍

图像分类是计算机视觉领域的基础任务之一，在智能安防、宠物管理、内容审核与教学实验等场景中具有广泛的应用价值。针对传统人工鉴别效率低、主观性强，以及从零训练深度模型对算力与样本量要求较高等问题，本文设计并实现了一套基于PyTorch的猫狗图像识别系统。系统以Python为主要开发语言，采用ResNet18卷积神经网络作为核心分类模型，并基于ImageNet预训练权重开展迁移学习；同时使用PyQt6构建桌面图形界面，实现模型训练、单张图像识别与数据集统计分析等功能的一体化集成。

在方法层面，系统通过冻结骨干网络、仅训练分类头的策略，显著降低CPU环境下的训练成本；结合随机裁剪、水平翻转与色彩抖动等数据增强手段，提升模型泛化能力；训练过程中实时绘制损失与准确率曲线，并自动保存验证集上表现最优的模型权重。实验结果表明，在每类采样约2000张图像、训练5个轮次的设置下，系统最优验证准确率可达99.00%，能够满足本科毕业设计对功能完整性、可演示性与技术深度的要求。

本文从需求分析、总体设计、详细实现到系统测试对课题进行了完整阐述，重点介绍了Python语言特性、PyTorch深度学习框架、ImageNet大规模数据集以及ResNet18残差网络等关键技术，并给出了各功能模块的核心代码说明。研究成果可为同类图像分类桌面应用的设计与教学实践提供参考。

源码下载

链接: https://pan.baidu.com/s/1z05iC6wDJmnsAYApQ0tXXQ?pwd=1234
提取码: 1234

系统展示

核心代码

""" 模型训练模块 使用 QThread 在后台执行训练，通过信号与 UI 通信 """ import json import os from typing import Optional import torch import torch.nn as nn import torch.optim as optim from PyQt6.QtCore import QThread, pyqtSignal import config from src.dataset import create_dataloaders from src.model import build_model, save_model from src.utils import format_datetime class TrainThread(QThread): """ 后台训练线程 在独立线程中执行模型训练，避免阻塞 UI """ # 信号定义 log_signal = pyqtSignal(str) # 日志消息 epoch_done_signal = pyqtSignal(dict) # 每轮训练完成 batch_progress_signal = pyqtSignal(int, int) # 批次进度 (current, total) finished_signal = pyqtSignal(bool, str) # 训练结束 (success, message) def __init__(self, parent=None): """ 初始化训练线程 :param parent: 父对象 """ super().__init__(parent) self._stop_flag = False # 训练参数 self.epochs = config.DEFAULT_EPOCHS self.batch_size = config.DEFAULT_BATCH_SIZE self.lr = config.DEFAULT_LR self.img_size = config.DEFAULT_IMG_SIZE self.subset_per_class = config.DEFAULT_SUBSET_PER_CLASS self.val_split = config.DEFAULT_VAL_SPLIT self.freeze_backbone = config.DEFAULT_FREEZE_BACKBONE def set_params( self, epochs: int = None, batch_size: int = None, lr: float = None, img_size: int = None, subset_per_class: int = None, val_split: float = None, freeze_backbone: bool = None, ): """ 设置训练超参数 :param epochs: 训练轮数 :param batch_size: 批次大小 :param lr: 学习率 :param img_size: 图像尺寸 :param subset_per_class: 每类子集数量 :param val_split: 验证集比例 :param freeze_backbone: 是否冻结骨干 """ if epochs is not None: self.epochs = epochs if batch_size is not None: self.batch_size = batch_size if lr is not None: self.lr = lr if img_size is not None: self.img_size = img_size if subset_per_class is not None: self.subset_per_class = subset_per_class if val_split is not None: self.val_split = val_split if freeze_backbone is not None: self.freeze_backbone = freeze_backbone def stop(self): """请求停止训练""" self._stop_flag = True self._emit_log("正在停止训练...") def _emit_log(self, message: str): """ 发送带时间戳的日志 :param message: 日志内容 """ timestamp = format_datetime() self.log_signal.emit(f"[{timestamp}] {message}") def run(self): """执行训练主流程""" try: self._stop_flag = False device = torch.device(config.DEVICE) self._emit_log(f"使用设备: {device}") self._emit_log(f"训练参数: epochs={self.epochs}, batch_size={self.batch_size}, " f"lr={self.lr}, subset={self.subset_per_class}/类") # 加载数据 self._emit_log("正在加载数据集...") train_loader, val_loader, dataset_info = create_dataloaders( batch_size=self.batch_size, img_size=self.img_size, subset_per_class=self.subset_per_class, val_split=self.val_split, ) self._emit_log(f"数据集加载完成: 训练集 {dataset_info['train_size']} 张, " f"验证集 {dataset_info['val_size']} 张") # 构建模型 self._emit_log("正在构建 ResNet18 模型...") model = build_model(num_classes=2, freeze_backbone=self.freeze_backbone) model.to(device) # 优化器与损失函数 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=self.lr ) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.5) # 训练历史 history = { "train_loss": [], "val_loss": [], "train_acc": [], "val_acc": [], } best_val_acc = 0.0 total_batches = len(train_loader) # 训练循环 for epoch in range(1, self.epochs + 1): if self._stop_flag: self._emit_log("训练已被用户停止") self.finished_signal.emit(False, "训练已停止") return self._emit_log(f"--- 第 {epoch}/{self.epochs} 轮 ---") # 训练阶段 model.train() train_loss, train_correct, train_total = 0.0, 0, 0 for batch_idx, (images, labels) in enumerate(train_loader): if self._stop_flag: self._emit_log("训练已被用户停止") self.finished_signal.emit(False, "训练已停止") return images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) train_correct += predicted.eq(labels).sum().item() train_total += labels.size(0) self.batch_progress_signal.emit(batch_idx + 1, total_batches) train_loss /= train_total train_acc = train_correct / train_total # 验证阶段 model.eval() val_loss, val_correct, val_total = 0.0, 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) val_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) val_correct += predicted.eq(labels).sum().item() val_total += labels.size(0) val_loss /= val_total val_acc = val_correct / val_total scheduler.step() # 记录历史 history["train_loss"].append(round(train_loss, 4)) history["val_loss"].append(round(val_loss, 4)) history["train_acc"].append(round(train_acc, 4)) history["val_acc"].append(round(val_acc, 4)) self._emit_log( f"Epoch {epoch}: train_loss={train_loss:.4f}, train_acc={train_acc:.2%}, " f"val_loss={val_loss:.4f}, val_acc={val_acc:.2%}" ) # 发送 epoch 完成信号 self.epoch_done_signal.emit({ "epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "train_acc": train_acc, "val_acc": val_acc, }) # 保存最优模型 if val_acc > best_val_acc: best_val_acc = val_acc save_model(model) self._emit_log(f"验证准确率提升至 {val_acc:.2%}，已保存最优模型") # 保存训练历史 with open(config.HISTORY_PATH, "w", encoding="utf-8") as f: json.dump(history, f, ensure_ascii=False, indent=2) self._emit_log(f"训练完成！最优验证准确率: {best_val_acc:.2%}") self.finished_signal.emit(True, f"训练完成，最优验证准确率: {best_val_acc:.2%}") except Exception as e: self._emit_log(f"训练出错: {str(e)}") self.finished_signal.emit(False, f"训练出错: {str(e)}")

""" 模型推理预测模块 加载训练好的模型，对单张图片进行猫/狗分类 """ import os from typing import Tuple, Optional import torch import torch.nn.functional as F from PIL import Image import config from src.model import load_model from src.dataset import get_predict_transform class Predictor: """ 猫狗图像分类预测器 封装模型加载与单图推理逻辑 """ def __init__(self, model_path: str = None): """ 初始化预测器 :param model_path: 模型权重路径 """ self.model_path = model_path or config.BEST_MODEL_PATH self.model = None self.transform = get_predict_transform(config.DEFAULT_IMG_SIZE) self.device = torch.device(config.DEVICE) self._loaded = False def load(self) -> bool: """ 加载模型 :return: 是否加载成功 """ try: self.model = load_model(self.model_path, num_classes=2) self._loaded = True return True except Exception as e: print(f"[预测] 模型加载失败: {e}") self._loaded = False return False def is_model_available(self) -> bool: """ 检查模型文件是否存在 :return: 模型是否可用 """ return os.path.exists(self.model_path) def predict(self, image_path: str) -> Tuple[str, float, dict]: """ 对单张图片进行预测 :param image_path: 图片路径 :return: (预测类别中文名, 置信度, 各类别概率字典) """ if not self._loaded: if not self.load(): raise RuntimeError("模型未加载，请先训练模型") if not os.path.exists(image_path): raise FileNotFoundError(f"图片不存在: {image_path}") # 加载并预处理图片 image = Image.open(image_path).convert("RGB") input_tensor = self.transform(image).unsqueeze(0).to(self.device) # 推理 self.model.eval() with torch.no_grad(): outputs = self.model(input_tensor) probabilities = F.softmax(outputs, dim=1)[0] # 解析结果 pred_idx = probabilities.argmax().item() confidence = probabilities[pred_idx].item() label = config.CLASS_NAMES.get(pred_idx, "未知") prob_dict = { config.CLASS_NAMES[i]: probabilities[i].item() for i in range(len(config.CLASS_NAMES)) } return label, confidence, prob_dict def predict_from_pil(self, image: Image.Image) -> Tuple[str, float, dict]: """ 对 PIL Image 对象进行预测 :param image: PIL Image 对象 :return: (预测类别中文名, 置信度, 各类别概率字典) """ if not self._loaded: if not self.load(): raise RuntimeError("模型未加载，请先训练模型") image = image.convert("RGB") input_tensor = self.transform(image).unsqueeze(0).to(self.device) self.model.eval() with torch.no_grad(): outputs = self.model(input_tensor) probabilities = F.softmax(outputs, dim=1)[0] pred_idx = probabilities.argmax().item() confidence = probabilities[pred_idx].item() label = config.CLASS_NAMES.get(pred_idx, "未知") prob_dict = { config.CLASS_NAMES[i]: probabilities[i].item() for i in range(len(config.CLASS_NAMES)) } return label, confidence, prob_dict