简介:面向图像分类与深度学习入门者的动物图片数据集,覆盖狗、猫、马、蜘蛛、蝴蝶、鸡、羊、牛、松鼠、大象10个类别,图片总量达28000余张,适合训练卷积神经网络、对比不同识别模型,也可模拟用户拍摄图像的真实分类场景。下载压缩包采用7z格式,内含2000个文件,主体为JPG/JPEG图片,另附8个PNG和1个Python脚本,包体大小约569.39MB;目录按类别分文件夹存放,便于快速导入训练或评估流程。作者使用该数据测试过自制CNN,准确率约80%,而迁移到谷歌Inception网络后可接近98%,能直观体现模型能力差异,也可为生物图像检索、智能画廊等应用提供数据基础。目前已有142人学习下载,是一套类别明确、来源经过人工检查的动物图像入门数据集。
1. 拿到 10 类动物 28000+ 张 JPG 数据集后,第一件事不是炼丹
做图像分类、目标检测或者迁移学习实验的人,迟早会碰到这种“10个不同类别动物图片数据集(28000+张图像)JPG”。它的形式通常是一个大文件夹,里面按类别分子目录,每张图是 JPG,总量两万八。听起来比 CIFAR-10 规模大、比 ImageNet 好啃,但真正上手会发现:这类数据集的坑不在“图不够多”,而在“没有统一的标签文件”“图片尺寸参差”“类别里混着模糊图和重复图”。如果你直接拿去跑分类训练,大概率前几个 epoch 就报错或者验证集指标虚高。
这篇文章不打算讲某个具体项目作者的代码,而是站在“拿到这种 JPG 数据集之后怎么把它变成能复现实验的训练管线”这个角度,把目录组织、数据体检、PyTorch 加载、YOLO 格式转换、常见踩坑一次讲清楚。适合正在做课程设计、公司内部算法验证、或者想用动物数据集练手图像分类/检测的工程师。目标就一个:让你拿到手之后,少走弯路,把时间花在调模型而不是修数据上。
2. 数据体检:28000+ 张 JPG 能不能直接训,先看三样东西
2.1 拿到数据集先看目录:10 类 JPG 的标准排列方式
这种“按类别分文件夹”的数据集,最常见的组织方式是:
animals10/ ├── cat/ │ ├── 00001.jpg │ ├── 00002.jpg │ └── ... ├── dog/ │ ├── 00001.jpg │ └── ... ├── bird/ └── ...我的建议是拿到手第一件事不是写训练代码,而是先把这个目录完整列出来看看。用几行命令确认有没有隐藏文件、空文件夹、非 JPG 文件混进来。
# 统计每个类别的文件数量 for d in animals10/*/; do cls=$(basename "$d") count=$(ls "$d" | wc -l) echo "$cls: $count" done这段脚本的逻辑很简单:遍历每个子目录,打印类名和文件数。跑完之后你心里就有数了——哪些类多、哪些类少,是不是长尾分布。注意看两点:一是有没有类别的文件数明显少于平均(比如只有 500 张,别的类 3000 张),二是文件名是否连续可预测。如果文件名是纯数字,后面做训练/验证划分时可以直接用哈希或者随机数,不需要重新命名。
接着再确认扩展名是否统一。很多声称“JPG”的数据集里实际混着 PNG、BMP,甚至一些无扩展名文件。用 find 过一遍:
find animals10 -type f | sed 's/.*\.//' | sort | uniq -c如果输出里出现 png、jpeg、gif 甚至没有扩展名的行,说明数据不干净。不用急着删,先记录下来,后面做数据加载时统一过滤,只在代码里接受.jpg和.jpeg。
2.2 用 Python 给 28000 张 JPG 做完整体检:坏图、尺寸与颜色空间
JPG 是压缩格式,网络上下载的数据集难免有截断文件、损坏文件、错误颜色空间。这些在图像查看器里可能只是“打不开”,但在 PyTorch 训练里就是直接的 RuntimeError。我的习惯是写一个独立的体检脚本,在训练之前跑完整个数据集,输出三类信息:打不开的图、异常尺寸的图、灰度图。
import os from PIL import Image from collections import Counter root = "animals10" bad_images = [] size_counter = Counter() gray_count = 0 for cls in os.listdir(root): cls_dir = os.path.join(root, cls) if not os.path.isdir(cls_dir): continue for fname in os.listdir(cls_dir): if not fname.lower().endswith((".jpg", ".jpeg")): continue path = os.path.join(cls_dir, fname) try: with Image.open(path) as img: img.verify() # 只校验完整性,不加载像素 w, h = img.size size_counter[(w, h)] += 1 except Exception: bad_images.append(path) print(f"坏图数量: {len(bad_images)}") for path in bad_images[:20]: print(path) print(f"不同尺寸组合数量: {len(size_counter)}") for size, cnt in size_counter.most_common(10): print(size, cnt)这里用img.verify()而不是img.load(),目的是只检查文件头与解码器是否匹配,不把整张图读进内存,28000 张跑起来很快。逻辑说明:错误会被 except 捕获并记录路径,后续你可以决定是删除还是用convert("RGB")容错。需要说明的是,verify()之后这张 Image 对象不能继续用,必须重新open()。
参数说明:size_counter用来统计尺寸分布,如果绝大多数图是类似 224x224、512x512 这样的尺寸,只有少数是异形尺寸,训练时统一 resize 就行;如果尺寸组合上百种,说明数据来源混杂,后面要重点考虑裁剪策略而不是简单 resize。灰度图在这个脚本里没有单独统计,你可以在读取时检查img.mode != "RGB"并打印,这类图如果不转 RGB,预训练模型输入通道对不上。
2.3 类别分布一眼看清:画分布图比看平均数有用
10 类 28000 张,按平均算每类 2800 张,但真实数据几乎不可能平均。有的类可能有 5000 张,有的只有 800 张。直接看数字可能会忽略长尾效应,画个柱状图最直观。
import os import matplotlib.pyplot as plt root = "animals10" classes = sorted([d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]) counts = [len(os.listdir(os.path.join(root, c))) for c in classes] plt.figure(figsize=(10, 5)) plt.bar(classes, counts) plt.xticks(rotation=45) plt.ylabel("image count") plt.title("class distribution") plt.tight_layout() plt.savefig("class_distribution.png")这段代码没有特别多的参数,核心是可视化类别分布,让你直观看到哪些类是“大头”、哪些是“稀缺类”。如果某个类少于 1000 张,在 28000 张的数据集里就是明显的长尾,后续训练要么做类别加权,要么做过采样,要么干脆对稀缺类做数据增强补偿。另外建议把counts按从大到小排序再画,排序后的图能更清楚地看到分布的衰减趋势。
提示:如果打印出的类别名是中文或者带空格,建议先统一重命名为英文小写加下划线,例如
wild_cat。这不是强迫症,而是避免后续 PyTorch Dataset、YOLO 配置、命令行参数里出现编码或空格解析问题。
3. 用 PyTorch 把这批 JPG 变成可训练数据:Dataset 写法与划分策略
3.1 自定义 Dataset:从文件夹到可迭代样本的最短路径
PyTorch 的torchvision.datasets.ImageFolder能直接处理“按类分子目录”的数据集,但它有个问题:它对每个 epoch 都重新遍历目录,IO 开销大,而且你没法方便地做缓存、过滤坏图或者返回额外信息。对于 28000 张图,直接用 ImageFolder 不是不行,但我更推荐写一个轻量自定义 Dataset,控制权在自己手里。
import os from PIL import Image from torch.utils.data import Dataset class AnimalDataset(Dataset): def __init__(self, root, transform=None, allowed_ext=(".jpg", ".jpeg")): self.root = root self.transform = transform self.classes = sorted( [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))] ) self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)} self.samples = [] for cls in self.classes: cls_dir = os.path.join(root, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith(allowed_ext): self.samples.append( (os.path.join(cls_dir, fname), self.class_to_idx[cls]) ) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label = self.samples[idx] img = Image.open(path).convert("RGB") if self.transform: img = self.transform(img) return img, label逻辑说明:初始化阶段就把所有样本路径和标签一次性扫进内存列表,训练时不再访问文件系统目录,只按索引取文件。class_to_idx用sorted排序,保证类名到数字标签的映射稳定;这个映射在训练和推理时都必须一致,否则就是“训练时猫是 0,推理时猫成 2”的灾难。
参数说明:allowed_ext控制哪些文件被当成图像,顺手把.jpeg也放进去,避免一个扩展名问题导致样本数少了百来张;convert("RGB")把灰度图和 RGBA 图统一转成三通道,不然模型输入通道数会报错。这个 Dataset 没有做坏图过滤,跑训练前建议先过一遍 2.2 的体检脚本,把坏图路径记下来,在这里直接用 set 过滤掉。
3.2 用 stratify 做分层划分:随机划分是新手最容易犯的错
28000 张图,很多人直接random.shuffle然后按比例切。这个做法在类别均衡时问题不大,但遇到长尾数据,可能验证集里稀缺类只有寥寥几张,训练集里反而很多,导致验证集指标忽高忽低。正确做法是按类别分层划分,保证训练集和验证集里每个类的比例一致。
from sklearn.model_selection import train_test_split paths = [s[0] for s in dataset.samples] labels = [s[1] for s in dataset.samples] train_paths, val_paths, train_labels, val_labels = train_test_split( paths, labels, test_size=0.2, stratify=labels, random_state=42, )参数说明:stratify=labels是核心,它让划分后的训练集和验证集中每个类别的样本占比与原数据集一致。test_size=0.2表示 28000 张里 22400 张训练、5600 张验证;random_state=42是为了复现,换数字也可以,但同一个实验里必须固定。train_test_split返回的是列表,你可以基于它重建两个子 Dataset 实例,或者直接写一个接受路径列表的 Dataset 封装。
这里有一个细节:划分之后建议把train_paths和val_paths存成.txt文件,每行一个路径。好处是下次训练不用重新划分,而且能保证“这次实验和上次实验用的是同一份验证集”,否则每次随机划分的结果不同,模型调参时你无法判断指标变化是来自模型改进还是数据划分运气。
3.3 数据增强参数怎么调:动物照片的 4 个关键点
动物数据集和 CIFAR-10、ImageNet 这类“目标居中”的数据集不太一样,动物照片往往背景复杂、目标占比小、姿态多变。数据增强如果直接抄 ImageNet 那套,容易把动物裁掉或者增强过度。我一般这样设置:
from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.6, 1.0), ratio=(0.75, 1.33)), transforms.RandomHorizontalFlip(p=0.5), transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])参数说明:RandomResizedCrop的scale=(0.6, 1.0)比 ImageNet 常用的(0.08, 1.0)保守很多。原因是动物占画面比例通常较大,裁太多会把头部或身体关键部位裁掉,模型学到的是“背景纹理”而不是“动物本体”。ratio=(0.75, 1.33)限制了裁剪框的宽高比范围,避免裁出过于细长的区域。ColorJitter的三个通道参数适度放大,因为动物照片经常在野外拍摄,光照变化剧烈,适当增强亮度对比度能提升泛化,但hue不要超过 0.1,否则毛色会发生不自然的偏移(比如狗变绿)。验证集不做随机增强,只用Resize(256)加CenterCrop(224),保证评估指标稳定。
注意:如果你打算之后用 ImageNet 预训练模型做迁移学习,Normalize 的 mean/std 必须保持
[0.485, 0.456, 0.406]这一组,不要自己统计数据集均值替换。替换后预训练权重照样能用,但收敛速度和最终精度通常会变差;只有从零训练才有必要计算自己数据集的均值方差。
4. 把动物数据集转成 YOLO 可训格式:分类与检测两条路
4.1 转成 YOLO 分类格式:不写一行训练代码就能跑
很多人手里的 JPG 数据集只有类别文件夹,没有 bbox 标注。这种情况下想用 YOLO 也不是不行——YOLO 官方支持图像分类任务,只需要把数据整理成train/class1/*.jpg这样的结构。做法是把上一节划分出的train_paths和val_paths复制到新目录里。
import os import shutil from sklearn.model_selection import train_test_split src_root = "animals10" yolo_root = "animals10_yolo" test_ratio = 0.2 random_seed = 42 classes = sorted([d for d in os.listdir(src_root) if os.path.isdir(os.path.join(src_root, d))]) for split in ["train", "val"]: for cls in classes: os.makedirs(os.path.join(yolo_root, split, cls), exist_ok=True) for cls in classes: cls_dir = os.path.join(src_root, cls) files = [f for f in os.listdir(cls_dir) if f.lower().endswith(".jpg")] train_files, val_files = train_test_split( files, test_size=test_ratio, random_state=random_seed ) for f in train_files: dst = os.path.join(yolo_root, "train", cls, f) if not os.path.exists(dst): shutil.copy(os.path.join(cls_dir, f), dst) for f in val_files: dst = os.path.join(yolo_root, "val", cls, f) if not os.path.exists(dst): shutil.copy(os.path.join(cls_dir, f), dst)逻辑说明:这段脚本在目标目录不存在的文件才复制,避免重复拷贝 28000 张图浪费磁盘空间;如果已经复制过一部分,断点续跑也有保障。划分用的random_seed=42和上一节的random_state=42保持同一套随机序列。
转完后用 YOLO 命令直接开训:
yolo classify train data=animals10_yolo model=yolov8n-cls.pt epochs=50 imgsz=224 batch=64参数说明:imgsz=224是输入分辨率,和 PyTorch 里RandomResizedCrop(224)保持一致,迁移学习时预训练权重不需要重新调整;batch=64取决于显存,24GB 显存可以尝试 128,8GB 建议 32 以下。yolov8n-cls.pt是分类版的轻量预训练模型,先拿它跑通全流程,验证数据没问题再换 s 或 m 版本。
这里有个判断标准:如果第一轮训练的 loss 在 5 个 epoch 内明显下降,说明数据链路是通的;如果 loss 纹丝不动或者直接 NaN,先回第 2 章检查数据。
4.2 想转目标检测?没有 bbox 标注时先别硬来
有的项目要求“识别动物位置”,这就得把分类数据集升级成检测数据集。问题是:标题给的数据集只有 JPG 和类别,没有标注框。常见的做法有两个:一是用标注工具(比如 LabelImg 或 X-AnyLabeling)人工框一部分,二是先用训练好的检测模型做预标注再人工修正。
如果你决定人工标注,YOLO 检测格式是这样组织的:
animals10_det/ ├── images/ │ ├── train/ │ │ ├── cat_00001.jpg │ │ └── ... │ └── val/ ├── labels/ │ ├── train/ │ │ ├── cat_00001.txt │ │ └── ... │ └── val/ └── classes.txtlabels/train/cat_00001.txt里每一行是一个目标,格式为:
class_id x_center y_center width height注意这四个坐标全部是归一化到 0~1 的浮点数,除以图片宽高。classes.txt每行一个类名,顺序必须和class_id一致。这个格式最大的坑是坐标类型:很多人直接把像素坐标写进去,训练时 loss 直接发散。
对 28000 张图全量标注不现实,我的建议是每个类别先标 200 张,总共 2000 张,用来验证检测管线能不能跑通。等模型在验证集上有基本效果,再用它去预测剩余图片并人工抽检修正,节省的时间是数量级的。
4.3 只有类别没有标注时,先用分类结果做“伪检测”
还有一种折中方案:不标 bbox,先跑分类模型,然后对每张图用 Grad-CAM 或注意力图观察模型关注区域。如果模型关注的是动物本体而不是背景,说明数据质量没问题;如果关注点跑到草地里,大概率是类别里混入了背景占比过大的图。
import torch from torchvision import models, transforms from PIL import Image model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) model.eval() img = Image.open("animals10/cat/00001.jpg").convert("RGB") transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) input_tensor = transform(img).unsqueeze(0) with torch.no_grad(): output = model(input_tensor)逻辑说明:这是标准的 ImageNet 预训练推理,不针对你的 10 类数据,目的只是快速让模型读一张 JPG 并输出预测。如果模型输出的类别和你期望的类别差异很大,不代表模型错了,更大概率是这张图本身的语义和类别标签不符——比如“猫”文件夹里混了一张狗的照片。这种跨类别混入在动物数据集里非常常见,尤其是相近物种(猫和豹猫、狗和狼)。
5. JPG 数据集训练避坑记录:五条血泪经验
5.1 现象:训练到一半报错 “image file is truncated”
训练 28000 张 JPG 时,常在某个 epoch 中途抛异常:
PIL.UnidentifiedImageError: cannot identify image file原因:下载或解压过程中部分 JPG 文件损坏,文件头还在但图像数据不完整。Image.open()不报错,但img.load()或 transform 触底解码时才炸。
解决:在 Dataset 初始化阶段就做一次全量校验,把坏图路径排除掉。
from PIL import Image, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True def is_valid_jpg(path): try: with Image.open(path) as img: img.load() return True except Exception: return False valid_samples = [ (p, l) for p, l in dataset.samples if is_valid_jpg(p) ] print(f"原始样本: {len(dataset.samples)},有效样本: {len(valid_samples)}")ImageFile.LOAD_TRUNCATED_IMAGES = True的作用是允许 PIL 解码截断文件,得到的图可能带色块,但至少不崩训练。我的建议是:训练集里遇到截断图直接丢弃,验证集里也丢弃,别依赖这个开关硬扛。因为带色块的图会污染模型学到的特征,坏图数量如果超过总样本的 1%,模型精度会明显下降。
5.2 现象:训练集准确率 97%,验证集 93%,测试集只有 70%
这种情况不是模型过拟合,而是数据划分泄漏。具体表现是:验证集和训练集里存在大量重复或近重复图像(比如同一只动物的连拍帧),random 划分把同一张图的不同版本分到了两边,模型其实在“背图”而不是“学特征”。
解决:先做重复图检测,再重新划分。
import hashlib from PIL import Image def image_hash(path, hash_size=8): img = Image.open(path).convert("RGB").resize((hash_size, hash_size)) # 简单的均值哈希 gray = img.convert("L") pixels = list(gray.getdata()) avg = sum(pixels) / len(pixels) return "".join(str(int(p > avg)) for p in pixels) hash_to_path = {} dup_groups = [] for path, _ in dataset.samples: h = image_hash(path) if h in hash_to_path: dup_groups.append((hash_to_path[h], path)) else: hash_to_path[h] = path print(f"疑似重复图对数: {len(dup_groups)}")逻辑说明:image_hash把每张图缩成 8x8 灰度图并计算均值哈希,相同哈希的图大概率是视觉近似的重复图。缩略图会把细节抹掉,所以这个方法的召回率较高但误报率也高——不同背景的两张猫照片可能算出相同的哈希。因此输出结果需要人工确认,确认后再决定从训练集里剔除其中一张。
5.3 现象:猫和老虎的类别始终互相混淆,loss 正常但准确率上不去
原因:数据集的类别定义本身有问题。有些“猫”的照片里实际包含老虎、豹子这类大型猫科动物,或者“狗”类别里有狐狸。对模型来说这些是不同物种,但对数据集标注来说它们是同一个文件夹。这类问题靠调模型解决不了,只能回到数据层面。
解决:检查混淆矩阵,找出互相错分的类别对,然后人工抽查这些类的图片。如果确认是标注混入,把错分的图片移动到正确类别,或者干脆删掉。不要试图用“增加训练轮数”来解决标注错误,那只会让模型更坚定地学习错误标签。如果错分情况集中在某一个人工标注来源,直接排查该来源的文件批次,用文件名前缀定位。
5.4 现象:Windows 下解压 zip 后目录乱码,或者有人问“jpg 文件怎么改成 zip”
先说结论:把 JPG 扩展名改成 .zip 并不能让图片变成压缩包,这只是骗过文件资源管理器,解压时反而得到一堆无法识别的二进制文件。真实场景是:你拿到的数据集是 zip 压缩包,在 Windows 上双击解压后,中文目录名变成乱码,导致代码里找不到类别文件夹。
原因:zip 内文件名编码是 GBK 或 UTF-8,Windows 自带解压工具按本地区码处理,遇到跨编码的压缩包就会乱码。
解决:用 Python 的 zipfile 解压并显式指定编码:
import zipfile zpath = "animals10.zip" out_dir = "animals10" with zipfile.ZipFile(zpath, "r") as zf: for name in zf.namelist(): # 处理中文乱码:尝试 UTF-8,失败则按 GBK 解码 try: decoded = name.encode("cp437").decode("utf-8") except UnicodeDecodeError: try: decoded = name.encode("cp437").decode("gbk") except UnicodeDecodeError: decoded = name zf.extract(name, out_dir)逻辑说明:zipfile 读取文件名时默认按 cp437 解码,导致原先是 UTF-8 或 GBK 的中文文件名变成乱码;这里先重新编码为 cp437 字节再按 UTF-8 解码,失败则按 GBK 试。解压后如果目录名仍然不是你想要的样子,检查一下decoded的输出内容再手动调整映射。这个问题的本质是文件名编码,和数据本身无关,但处理不好会浪费半天。
注意:如果数据集是 Linux 上打包的,文件名大多已经是 UTF-8;Windows 上打包的则可能带 GBK。所以解压前先看一眼 zip 列表,比解压后猜编码强得多。
5.5 现象:28000 张 JPG 加载慢到训练一个 epoch 要 20 分钟
原因:每张图都是几百 KB 到几 MB 的 JPG,训练时实时解码,CPU 成为瓶颈。尤其是 Resize 到 224 也需要在 CPU 上完成,GPU 在那里空转等你。
解决:预处理缓存。第一次训练前把所有图统一 resize 到目标尺寸并转成 tensor 保存到磁盘,之后每次训练直接加载缓存文件。一个简单的做法是把所有图像缓存成一个.pt文件:
import torch from PIL import Image def preprocess_and_cache(dataset, output_path, size=224): tensors = [] labels = [] transform = transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), ]) for i in range(len(dataset)): img, label = dataset[i] tensors.append(transform(img)) labels.append(label) data = { "tensors": torch.stack(tensors), "labels": torch.tensor(labels), } torch.save(data, output_path) print(f"缓存完成,共 {len(tensors)} 张图")逻辑说明:缓存一次后,训练时不再做 JPG 解码和 resize,直接把整个 tensor 加载进内存。28000 张 224x224 的 RGB tensor 大约占用 28000 * 3 * 224 * 224 * 4 字节,约 16GB 内存,内存充足的情况下这个方案最省事。如果内存不够,按类别拆成多个缓存文件,或者保存成 PNG 而不是 JPG——PNG 解码快但不压缩模型训练强度。
参数说明:size=224必须和训练时的输入尺寸一致,否则缓存后还需要再次 resize,等于白缓存。如果你之后要试 256 或 320 分辨率,缓存就得重新生成,所以缓存尺寸最好定成你最终的训练尺寸,不要反复变。
6. 训练前最后一道验证:跑通冒烟测试,再看混淆矩阵
6.1 先跑一个 step 和一个 10 epoch 的冒烟测试
正式训练前,我习惯做两件事:第一,从 Dataset 里取一个 batch,检查 tensor 形状和标签范围;第二,用极小模型跑 10 个 epoch,确认 loss 能从高值降下来。
from torch.utils.data import DataLoader train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4) batch, labels = next(iter(train_loader)) print("batch shape:", batch.shape) # 期望 torch.Size([32, 3, 224, 224]) print("labels:", labels.unique()) # 期望包含 0~9 的数字,且没有缺失 # 用一个小模型做冒烟测试 import torch.nn as nn model = nn.Sequential( nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(16, 10), ) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for epoch in range(10): total_loss = 0 for i, (imgs, lbls) in enumerate(train_loader): out = model(imgs) loss = criterion(out, lbls) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"epoch {epoch}, loss: {total_loss / len(train_loader):.4f}")batch.shape如果不对,先检查 transform 里有没有ToTensor,常见错误是忘了加这一层导致输入是 PIL Image 而不是 tensor。labels.unique()如果缺少某个类,说明该类的样本数为 0,返回去查目录结构。冒烟测试的 10 个 epoch 里,loss 应该从 2.3 左右(10 类的随机交叉熵)逐步降到 1.5 以下;如果 loss 涨到 10 以上或者变成 NaN,优先怀疑学习率太大,其次怀疑数据里有异常的像素值——用imgs.max()和imgs.min()检查一下,Normalize 后最大值不应超过 3 个标准差太多。
6.2 混淆矩阵比 accuracy 更值得看
训练完成后,我看的第一张图永远是混淆矩阵,而不是 accuracy。因为 accuracy 只告诉你好不好,混淆矩阵告诉你哪些类别互相打架。
import numpy as np from sklearn.metrics import confusion_matrix def eval_confusion(model, val_loader, classes): all_preds = [] all_labels = [] model.eval() with torch.no_grad(): for imgs, lbls in val_loader: out = model(imgs) preds = out.argmax(dim=1) all_preds.extend(preds.tolist()) all_labels.extend(lbls.tolist()) cm = confusion_matrix(all_labels, all_preds) np.fill_diagonal(cm, 0) # 只看错误 for i, cls in enumerate(classes): for j in range(len(classes)): if cm[i][j] > 0: print(f"{classes[i]} -> {classes[j]}: {cm[i][j]}")逻辑说明:np.fill_diagonal(cm, 0)把正确预测清零,剩下的非零项全是错分。输出格式是“真实类 -> 预测类”,按数量排序就能快速定位问题最大的类别对。如果某个非对角元素特别大,回到 5.3 检查类别定义;如果是相邻形态物种互串(比如猫和虎),考虑增加细粒度分类的数据增强或者干脆合并类别。
这个验证习惯我一直保留着:拿到手的数据先体检再炼丹,训练前跑冒烟测试,训练后先看混淆矩阵再看 accuracy。这套流程看起来慢,实际是避免返工最快的方式。28000 张 JPG 的数据集不算小,数据没准备好就开训练,等 50 个 epoch 跑完发现是数据问题,那才是真的血泪教训。希望帮到你。
本文还有配套的精品资源,点击获取