简介基于PyTorch与ShuffleNet模型的宠物年龄识别代码包适合Python/CNN学习者完成图片分类实战。项目预置完整工作流01生成txt.py将图片路径与标签转为训练/验证集文本02CNN训练数据集.py自动读取txt并适配任意分类文件夹03pyqt界面.py加载model.ckpt进行图片识别训练过程显示进度条、准确率与损失值并自动保存log日志。资源共9个文件含3个Python脚本、1个说明文档、1个依赖清单、4张示意jpg压缩包仅228KB轻量易用。目前已有112人学习适合需要快速搭建宠物年龄识别训练流程的开发者。下载后自行放置图片即可直接训练逐行注释与说明文档可辅助理解关键代码便于二次开发与课程设计参考。1. 轻量模型判断宠物年龄为什么选了 ShuffleNet 而不是 ResNet宠物年龄识别这个任务真正难的不是区分猫狗而是同一物种在不同年龄段的外观差异太小——幼猫和成猫差距往往只在脸型和骨架比例上。拿 ResNet 这类重网络去拟合小规模宠物照片参数量大、推理慢在小数据集上还容易过拟合动不动就在训练集上跑到 99% 准确率换一张真实照片立刻现原形。这个项目选择 ShuffleNet 解决的是实际落地场景用普通 CPU 做推理训练脚本只有三个数据准备、模型训练、PyQt 界面识别各管一段代码带逐行注释和说明文档适合课程设计、毕业设计也适合想在轻量网络上快速跑通完整图像分类流程的 Python 开发者。整个工程不含数据集图片需要自己按类别整理照片这会逼着你把深度学习的数据组织逻辑弄清楚。2. 数据集组织与标签生成01生成txt.py 的目录扫描逻辑2.1 训练集目录结构与放置规范这个项目的数据输入约定很直白每个年龄类别一个文件夹文件夹名就是类别名。以压缩包里的四个类别为例常见做法是组织成下面这种结构dataset/ ├── 幼狗/ │ ├── 幼狗 1.jpg │ └── ... ├── 成狗/ │ ├── 成狗 1.jpg │ └── ... ├── 幼猫/ │ ├── 幼猫 1.jpg │ └── ... └── 成猫/ ├── 成猫 1.jpg └── ...需要特别注意的是压缩包里每个文件夹内都有一张提示图片提示你应该把搜集到的图片放到哪个位置。这个设计对新手友好但也埋了一个坑——运行 01生成txt.py 时如果扫描脚本不做过滤提示图本身会被当成训练样本写进 txt训练时模型会学到一个固定窗口降低真实场景下的泛化能力。处理方式有两种要么把提示图移走要么在扫描逻辑里对文件扩展名做白名单过滤。后面的代码采用第二种做法更稳妥。过滤项建议值原因图片扩展名.jpg .jpeg .png .bmp大小写不敏感排除提示图、隐藏文件和临时文件最小文件大小建议大于 5 KB跳过损坏或不完整的图片文件夹名编码保持 UTF-8避免 Windows 下中文目录名乱码导致标签映射错乱2.2 01生成txt.py 如何把目录转成训练清单这个脚本的核心工作是扫描数据目录给每个类别分配数字标签然后按比例切分训练集和验证集把路径与标签写成 txt。一个典型实现如下import os import random def generate_txt(root_dir, output_train, output_val, train_ratio0.8): # root_dir 为数据集根目录例如 ./dataset classes sorted([d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]) # 按字典序给每个分类文件夹分配数字标签保证多次运行结果一致 class_to_idx {cls: i for i, cls in enumerate(classes)} train_lines [] val_lines [] for cls in classes: cls_dir os.path.join(root_dir, cls) imgs [] for name in os.listdir(cls_dir): # 只保留图片文件避免把提示图/说明文档扫进来 if name.lower().endswith((.jpg, .jpeg, .png, .bmp)): imgs.append(os.path.join(cls_dir, name)) random.shuffle(imgs) # 先打乱再划分避免同类图片连续扎堆 split int(len(imgs) * train_ratio) for path in imgs[:split]: train_lines.append(f{path} {class_to_idx[cls]}\n) for path in imgs[split:]: val_lines.append(f{path} {class_to_idx[cls]}\n) with open(output_train, w, encodingutf-8) as f: f.writelines(train_lines) with open(output_val, w, encodingutf-8) as f: f.writelines(val_lines) print(ftrain: {len(train_lines)} 张, val: {len(val_lines)} 张)逻辑并不复杂先扫描 root_dir 下的第一级子目录目录名就是类别名按字母排序后分配从 0 开始的数字标签。随后在每个类别目录里筛出图片文件按 train_ratio 切分。输出的 txt 每行格式是图片绝对路径 数字标签中间用空格隔开这是后续训练脚本读取的标准格式。参数层面有几个值得注意的地方。train_ratio 默认 0.8如果你的数据量少于每类 50 张我建议调到 0.9否则验证集可能只有几张图评估出来的准确率波动非常大。存路径时尽量写绝对路径避免训练脚本的工作目录与数据集目录不一致时出现 FileNotFoundError。这里用 txt 而不是 CSV是因为 txt 的解析最省事也方便直接打开看一眼有没有扫入奇怪的文件。2.3 数据准备阶段就要想好的 transform 规划transform 的代码写在第 02 个训练脚本里但数据准备阶段必须想清楚。这个项目最终要喂给 ShuffleNet 的图片尺寸固定在 224x224因为 torchvision 里 ShuffleNet 的 ImageNet 预训练权重就是按 224x224 输入的。如果图片尺寸不一致模型 forward 时张量形状与全连接层维度对不上报错都是维度不匹配。训练集的 transform 建议包含随机水平翻转和随机裁剪验证集只做 resize 和中心裁剪这能让训练集有明显的数据扰动验证集保持稳定。Normalize 的均值方差直接用 ImageNet 的统计值即可。灰度图也要在这个阶段处理掉要么转成 RGB 三通道要么在读取时统一调用 convert(RGB)否则单通道输入会在卷积层直接报错。3. 训练脚本核心解析02CNN训练数据集.py 的类别适配与训练流程3.1 自定义 Dataset 从 txt 读取样本训练脚本第一步是定义 Dataset 类。为什么不用 torchvision 自带的 ImageFolder因为这个项目已经通过 txt 清单把数据索引固化了训练时不依赖目录结构后续新增类别或调整数据划分都不需要改训练脚本。from torch.utils.data import Dataset from PIL import Image class PetAgeDataset(Dataset): def __init__(self, txt_path, transformNone): self.samples [] self.transform transform self.classes set() # 收集出现过的所有标签 with open(txt_path, r, encodingutf-8) as f: for line in f: line line.strip() if not line: continue path, label line.split() self.samples.append((path, int(label))) self.classes.add(int(label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) # 统一转 RGB规避灰度图与 PNG 透明通道问题 if self.transform: img self.transform(img) return img, label这段代码有两个容易被忽略的细节。第一是 convert(RGB)宠物照片来源多样手机拍的可能带 EXIF 方向信息网上下载的可能是灰度图或带 alpha 通道的 PNG如果不统一通道数训练中途会因为某张图张量形状不一致而崩溃。第二是 classes 集合虽然只读最大值也能推出类别数但显式收集 set 更稳妥因为你要保证标签是连续的 0 到 N-1如果 txt 里出现标签缺失直接取 max1 会出错。3.2 动态类别数如何替换 ShuffleNet 的最后一层项目描述里强调适配了数据集的分类文件夹个数即使增加了分类文件夹也不需要修改代码即可训练。实现思路就是先统计数据集的类别数再根据类别数动态构造输出层。import torchvision.models as models import torch.nn as nn def build_model(num_classes, pretrainedTrue): # x0_5 是最轻量档位显存占用小如果显存充足可以换成 x1_0 model models.shufflenet_v2_x0_5( weightsmodels.ShuffleNet_V2_X0_5_Weights.IMAGENET1K_V1 if pretrained else None ) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) return model关键在于 model.fc.in_features 这个属性它返回上一层的输出维度。为什么不直接硬编码因为 ShuffleNet x0_5 和 x1_0 的 fc 输入维度不同不同版本 torchvision 内部实现也有差异硬编码会把自己绑死。替换后模型的卷积骨干保持不变只有最后输出的分类数量变化。关于 ShuffleNet 的结构需要多说两句。ShuffleNet v2 的基础块把输入按通道拆成两半一半直接走深度可分离卷积另一半做恒等映射最后拼接起来再用 channel shuffle 把通道顺序打乱。这个设计的直接收益是在相同 FLOPs 下精度比 v1 高实际推理速度也更快因为深度卷积和逐点卷积的访存开销被控制住了。对宠物年龄识别这种图像细节差异不大的细粒度任务它比 MobileNet 更擅长捕捉局部纹理变化。pretrained 参数的选择要看数据量。每类几百张这种规模加载 ImageNet 预训练权重只微调最后两层是最稳的方案。每类有上万张再考虑从零训练否则收敛速度慢且容易落在次优解上。3.3 训练循环进度条、损失、准确率与模型保存训练循环是整个项目的核心。它要把 DataLoader 里的数据一批批喂给模型每个 epoch 结束后算准确率和损失写日志并保存最新的权重。import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model build_model(num_classes).to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) for epoch in range(1, epochs 1): model.train() running_loss 0.0 correct 0 total 0 pbar tqdm(train_loader, descfEpoch {epoch}) for images, labels in pbar: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() pbar.set_postfix(lossf{loss.item():.4f}) epoch_loss running_loss / total epoch_acc correct / total with open(train_log.txt, a) as f: f.write(fepoch {epoch} loss {epoch_loss:.4f} acc {epoch_acc:.4f}\n) torch.save(model.state_dict(), model.ckpt)几个参数需要认真对待。学习率设成 1e-4 而不是常见的 1e-3因为加载了预训练权重后模型已经处在一个较优的局部区域学习率太大会把学到的通用特征冲掉。batch_size 建议 16 或 32ShuffleNet x0_5 很轻量4GB 显存也能跑 32CPU 训练耐心一点也能完成。scheduler 每 10 个 epoch 把学习率减半这样 loss 到达平台期后还能继续下降。CrossEntropyLoss 内部已经包含 softmax 操作网络输出层不需要额外加 softmax。如果需要看置信度在推理阶段用 torch.softmax 处理输出训练阶段直接拿原始 logits 算损失即可。超参数推荐值说明输入尺寸224x224与 ImageNet 预训练输入对齐batch_size16 / 32小显存用 16CPU 也建议 16学习率1e-4预训练模型微调常用范围epochs30~50观察 loss 平台期优化器Adam收敛快对新手友好损失函数CrossEntropyLoss多分类标准选择这里的 torch.save 保存的是 model.state_dict()只有权重没有结构。如果训练数据类别变了模型结构就变了旧权重没法直接加载。更稳妥的保存方式是存一个 dict把 model.state_dict、类别列表、输入尺寸一起打包后面第 4 章的推理脚本会用到这个结构。4. PyQt 推理界面03pyqt界面.py 从加载权重到输出预测4.1 加载 model.ckpt 的三个一致性检查推理脚本最容易阴沟翻船的地方不在 PyQt 布局而在数据预处理与训练时不一致。第一是尺寸训练用 224x224推理也必须先 resize 到 224x224否则卷积层接收到的张量空间尺寸不同直接报维度错误。第二是 Normalize 参数mean 和 std 两边必须一样顺序也不能反。第三是通道顺序PyTorch 要 CHW通道在前OpenCV 读出来是 HWC 且是 BGR所以尽量用 PIL 读图或者用 OpenCV 读完后显式转换。import torch from PIL import Image import torchvision.transforms as transforms def preprocess(image_path): transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) return transform(img).unsqueeze(0) # 加 batch 维 def predict(model, image_path, device, class_names): model.eval() tensor preprocess(image_path).to(device) with torch.no_grad(): output model(tensor) prob torch.softmax(output, dim1) conf, idx torch.max(prob, 1) return class_names[idx.item()], conf.item()predict 函数返回类别名和置信度。class_names 是一个 Python 列表顺序必须与训练时目录名的排序一致。这里 torch.no_grad() 是推理的基本规范不计算梯度省内存速度也更快。4.2 模型加载map_location 与结构重建前面训练时保存的是裸的 state_dict推理脚本加载时要先重建同样的模型结构再填充权重。如果训练时保存了包含 class_names 的完整 checkpoint加载端就省事得多。def load_model(ckpt_path, device): ckpt torch.load(ckpt_path, map_locationdevice) class_names ckpt[class_names] model build_model(len(class_names), pretrainedFalse) model.load_state_dict(ckpt[state_dict]) model.to(device).eval() return model, class_namesmap_location 这个参数很实用。训练在 GPU 上进行推理在 CPU 电脑上运行时不写 map_location 会报 Attempting to deserialize object on a CUDA device 的错。统一写成 device 变量加载时 torch 会按目标设备自动处理。加载完调用 model.eval() 是必须的它会关闭 Dropout 和 BatchNorm 的训练模式保证推理结果稳定。如果只有裸的 state_dict也可以通过扫描模型文件的 fc 输出维度来反推类别数或者用第 2 章生成的 txt 重新读取类别列表但都不如保存 checkpoint 时附带 class_names 来得干净。4.3 PyQt 界面交互逻辑与卡顿规避界面逻辑不复杂一个按钮选择图片一个按钮触发推理两个 QLabel 分别显示预览图和结果。关键点在槽函数里不要做重活。from PyQt5.QtWidgets import QFileDialog, QLabel, QPushButton def choose_image(self): path, _ QFileDialog.getOpenFileName( self, 选择图片, , 图片文件 (*.jpg *.jpeg *.png *.bmp)) if path: self.image_path path self.preview_label.setPixmap(QPixmap(path).scaled(400, 400)) def run_predict(self): if not hasattr(self, image_path): return name, conf predict(self.model, self.image_path, self.device, self.class_names) self.result_label.setText( f预测结果{name}置信度{conf:.2%})对 ShuffleNet x0_5 来说CPU 上单张 224x224 图片的推理通常在 1 秒内放在按钮槽函数里不会明显卡顿。但如果以后换成 ResNet 或批量预测就要把推理放到 QThread 子线程里否则界面会冻结。另一个值得做的是把模型加载放到启动界面之前避免点击按钮时卡住。5. 验证模型到底学到了什么一张误判表定位分类瓶颈模型训练完别急着直接上手先把验证集完整跑一遍生成混淆矩阵和逐类别准确率。这里有一个很实用的排查脚本import numpy as np num_classes len(class_names) cm np.zeros((num_classes, num_classes), dtypeint) wrong_samples [] for path, true_label in val_samples: pred_name, conf predict(model, path, device, class_names) pred_label class_names.index(pred_name) cm[true_label][pred_label] 1 if pred_label ! true_label: wrong_samples.append((path, true_label, pred_label, conf)) for i, name in enumerate(class_names): total cm[i].sum() acc cm[i][i] / total if total else 0 print(f{name}: {acc:.2%}) wrong_samples.sort(keylambda x: -x[3]) # 按置信度降序 for path, true_l, pred_l, conf in wrong_samples[:10]: print(f{path} 真实{class_names[true_l]} 预测{class_names[pred_l]} 置信度{conf:.2%})这个脚本的价值在于把感觉模型还行变成具体哪类有问题。如果成猫和幼猫之间互相误判多先怀疑是数据混入了错误标签打开误判图片看一眼就知道。如果某个类别准确率明显低于其他优先补这一类的图片量比盲目调超参数有效得多。更进一步的分析手段是看热力图。取模型最后一个卷积层的输出特征图按类别得分对特征图做加权求和然后上采样到原图尺寸叠加显示就能定位模型在判断时关注的是宠物脸部还是背景。对年龄识别这种细粒度任务好的热力图应该集中在五官、毛发纹理和体型轮廓上如果高亮区域大面积落在背景上说明训练集里不同类别的拍摄场景太相似模型在用背景做捷径这时要把精力花在采集多样化场景的图片而不是继续调学习率。另外一个容易被忽略的筛选逻辑高置信度的误判样本优先级最高。置信度高但判错说明模型对这个样本非常确定这往往是标注错误或样本本身有歧义。置信度低判错的样本说明模型本来就在边界上犹豫改进价值反而不大。所以按置信度降序排列误判样本一条条排查是把有限时间花在刀刃上的做法。本文还有配套的精品资源点击获取