简介:面向图像识别与农业智能化应用的苹果品种分类数据集,适合机器学习、深度学习研究者及开发者用于训练苹果品种识别与分类模型,可服务于智能农业、食品质检、电商分拣等场景。资源包为zip压缩包,共1766个文件,包含580张高清苹果图像(305个jpg和275个jpeg),并配有580个xml标注文件与606个txt说明文档,整包约64.01MB,文件组织便于按品种读取和划分训练集。目前已有207人浏览学习。数据集覆盖多种苹果品种,图片包含全貌、切面、横截面等多角度细节,标注信息涉及品种名称、产地、口味特点等;txt与xml文件配合图像可提供类别标签与标注信息,可直接用于卷积神经网络等深度学习模型的训练与评估。配合多样化的图像背景,可帮助使用者快速搭建分类实验,验证算法在不同环境下的鲁棒性。
1. 苹果品种分类数据集:拿到 580 张 JPEG 之后的第一件事
做图像分类的同行第一次接触这份苹果品种分类数据集时,最常见的反应是直接解压、随手挑几张图看看、然后塞进训练脚本。这个流程我用过,翻车概率很高。这份数据集一共有 580 张高清 JPEG 图片,覆盖多个苹果品种,包含全貌、切面、横截面等角度,背景和拍摄条件并不统一——听起来像一份正经的分类数据集,但它不是 Kaggle 上下载好、带干净 CSV 标注的那种开箱即用的资源,文件名是203.jpeg、530.jpeg这样的流水号,品种信息并不会直接出现在文件名里,需要先做一轮数据盘点。这篇文章就围绕「怎么把这 580 张散装图片变成一套能训练、能评估、能复现的分类数据集」展开,适合做智能农业视觉识别、食品分级检测的算法工程师,也适合拿图像分类当毕业设计、想走完一遍完整 pipeline 的学生。
2. 摸清这份苹果数据集的底细:文件组织与快速预审
2.1 解压之后的目录里到底有什么
我通常的习惯是先把 zip 包解压到一个纯英文路径下,避免后面 PyTorch 的 DataLoader 在 Windows 上因为中文路径闹脾气。常见做法是这样:
mkdir -p ~/datasets/apple_classification && cd ~/datasets/apple_classification unzip ../apple_品种分类数据集.zip ls -la | head -20解压后你会看到一堆类似203.jpeg、530.jpeg、18.jpeg的图片文件,没有任何子目录。ls -la之后先数一下文件数量对不对得上简介里说的 580 张:
find . -name "*.jpeg" -o -name "*.jpg" | wc -l这一步看似多余,但很有必要——网络上下载的压缩包经常出现文件缺失或者解压中断的情况,先确认数量是 580 再继续。find命令里加了-o(or 的意思),把.jpeg和.jpg两种后缀都统计进来,实际这份数据集里以.jpeg为主,但保险起见我把两类后缀都扫了一遍。
数量没问题之后,我会顺手看一眼文件大小的分布。图片尺寸和体积直接决定了后续训练时要不要做统一缩放,也决定了数据加载的 IO 压力:
ls -la *.jpeg | awk '{print $5}' | sort -n | uniq -c | head -20这条命令把每个文件的大小打印出来、排序、去重计数,结果里能看到图片体积集中在哪个区间。如果大部分图片在 100KB 到 500KB 之间,属于正常手机或相机拍摄的 JPEG;如果出现大量 1MB 以上的大图,训练前就必须 Resize,否则显存会吃紧。
2.2 快速盲审:用脚本扫一遍图片的基础质量
文件数量对上了,不代表图片都能用。JPEG 在压缩包里传输过程中偶尔会损坏,PIL 能打开但里面是一张灰图或者花屏的情况也不少见。我一般会写一个扫描脚本,把打不开、模式异常、尺寸异常的图片全部挑出来:
from PIL import Image import os img_dir = "." bad_files = [] for name in sorted(os.listdir(img_dir)): if not name.lower().endswith((".jpeg", ".jpg")): continue path = os.path.join(img_dir, name) try: with Image.open(path) as im: im.load() w, h = im.size if w < 100 or h < 100: bad_files.append((name, "too_small", f"{w}x{h}")) if im.mode not in ("RGB", "RGBA"): bad_files.append((name, "bad_mode", im.mode)) except Exception as e: bad_files.append((name, "corrupt", str(e))) print(f"total bad files: {len(bad_files)}") for item in bad_files: print(item)这段脚本的逻辑分三层:第一层用im.load()真正把像素数据读进内存,这一步能过滤掉大部分头部信息完整但数据损坏的假图片;第二层检查尺寸,小于 100×100 的图片对分类任务几乎没用,直接标记;第三层检查色彩模式,灰度图、调色板图、CMYK 图都要单独记一笔,因为后续训练的输入层默认是三通道 RGB,提前知道有哪些模式异常能省掉运行时才爆出来的预处理错误。
参数上值得说明的是im.load()。很多人在做数据预审时只用Image.open()拿个尺寸就完事,但open()是惰性的,文件头能读、像素数据损坏时它不会报错,等训练中途__getitem__里再读取才炸,到时排查成本高得多。预审阶段宁可慢一点,把每一张图都完整加载一遍。
2.3 从文件名和像素信息反推出标签线索
扫描完基础质量,接下来要回答一个关键问题:这 580 张图分别属于哪些品种?上面说了,文件名是纯数字流水号,没有品种信息。这时有两个线索可以挖:一是图片的 EXIF 信息,二是图片本身的内容特征。
from PIL import Image from PIL.ExifTags import TAGS sample_files = ["203.jpeg", "530.jpeg", "18.jpeg", "626.jpeg"] for name in sample_files: with Image.open(name) as im: exif = im._getexif() if exif is None: print(f"{name}: no exif") continue for tag_id, value in exif.items(): tag_name = TAGS.get(tag_id, tag_id) if tag_name in ("Make", "Model", "DateTimeOriginal", "Software"): print(f"{name}: {tag_name} = {value}")这个脚本打印拍摄设备、拍摄时间、处理软件等信息。我的经验是,这类数据集里的 EXIF 大多数已经被清洗掉了,能挖到拍摄设备信息就算运气不错——但这至少能帮你判断图片是手机拍的还是相机拍的,不同设备的色彩倾向不同,后续做归一化时有参考价值。
EXIF 没有线索的情况下,最直接的办法是把全部 580 张图片拼成一张大缩略图,肉眼过一遍,把明显属于同一品种的图片先归堆。这一步听起来原始,但非常有效。我一般这么干:
from PIL import Image import os import math files = [f for f in sorted(os.listdir(".")) if f.lower().endswith((".jpeg", ".jpg"))] thumb_size = 64 cols = 20 rows = math.ceil(len(files) / cols) sheet = Image.new("RGB", (cols * thumb_size, rows * thumb_size), "white") for idx, name in enumerate(files): im = Image.open(name).resize((thumb_size, thumb_size)) x = (idx % cols) * thumb_size y = (idx // cols) * thumb_size sheet.paste(im, (x, y)) sheet.save("contact_sheet.jpg") print(f"saved contact_sheet.jpg with {len(files)} images")缩略图拼版的核心参数是thumb_size = 64和cols = 20。64 像素足够看出苹果的大致颜色和形状,20 列一行的排布在 1280 像素宽的屏幕上正好铺满,不用横向滚动。20 行一共能排 400 张,580 张图需要 29 行,生成的contact_sheet.jpg大概 1800 像素高,一屏看不完就上下滚动着看。如果你用的是 4K 屏幕,可以把thumb_size调到 96,看得更清楚。
这一步之后,你心里应该有一个初步的品种数量判断了:这份数据集里到底有 5 个品种还是 8 个品种、哪些品种的样本特别多、哪些品种只有零星几张——这个信息直接决定了第 4 章训练时的采样策略和损失函数要不要做类别加权。
3. 把原始 JPEG 整理成可训练的苹果分类数据集:标签提取与目录划分
3.1 标签体系怎么建:没有标注文件时的选择
数据预审做完了,接下来面对的是这份数据集最棘手的地方:没有 CSV,没有 JSON,没有 XML,任何形式的标注文件都不存在。品种信息藏在图片本身里,需要人工打标。
这时候有两种路径。第一种是逐张看图、人工归类,把203.jpeg记为red_delicious、530.jpeg记为granny_smith,逐张写映射关系。第二种是先按视觉相似度把 580 张图分成若干堆,再为每一堆分配一个品种名。我强烈建议用第二种,原因很简单:先粗后精,粗分时注意力集中在「像不像」,细分类时再核对「是不是同一个品种」,比逐张命名要快得多,而且不容易因为前后标准不一致造成两张同一品种的图被标成不同标签的脏数据。
在建立标签体系之前,还有一个问题要想清楚:类别粒度。苹果的分类可以细到品种名,比如富士、嘎啦、红蛇果、青香蕉;也可以粗到颜色或产地属性,比如红苹果、青苹果、黄苹果。这份数据集的描述里提到「每张图片均代表了不同的苹果品种」,说明原始采集意图是品种级分类,所以标签建议按品种名走。但如果你在盲审时发现某些品种的样本量差异过大——比如富士 200 张、某小众品种只有 8 张——就要考虑做类别合并,把特征相似的品种归并成一个类,否则训练时模型会对小样本类别严重欠拟合。
我一般会把标签映射放在一个独立的 Python 文件里,而不是直接建目录。原因很实际:后面划分训练集、验证集、测试集时要反复读标签,如果一开始就按目录结构移动文件,改标签的成本极高——你得把图片从一个目录移动到另一个目录,移动完还得记得更新清单。用映射文件的方式,图片始终待在原地,目录结构只是训练时的一种视图。
# label_map.py APPLE_LABELS = { "203.jpeg": "fuji", "530.jpeg": "granny_smith", "18.jpeg": "fuji", "626.jpeg": "red_delicious", "614.jpeg": "gala", "267.jpeg": "fuji", # 其余图片按同一方式补充,约 580 条 }写这个映射文件时注意,key 必须和文件名完全一致,包括大小写。Windows 文件系统大小写不敏感,Linux 和 macOS 默认大小写敏感,同一个203.JPEG在不同系统上可能读到不同文件。我的习惯是写一个校验脚本,把磁盘上实际存在的文件名和映射表的 key 做差集,避免训练时 KeyError 才发现少标了几十张。
3.2 按 ImageFolder 规范重组目录结构
标签映射建好之后,就可以把它转换成 PyTorch 标准的数据集目录结构了。PyTorch 的torchvision.datasets.ImageFolder约定:根目录下按类别分子目录,每个子目录名就是类别名,目录里放该类别的所有图片。这个结构的好处是不用自己写 Dataset 类,ImageFolder会自动扫描目录并生成索引,配合DataLoader就能直接进入训练流程。
import os import shutil from label_map import APPLE_LABELS src_dir = "." dst_dir = "apple_dataset" # 收集标签中出现的所有类别名 categories = sorted(set(APPLE_LABELS.values())) for c in categories: os.makedirs(os.path.join(dst_dir, c), exist_ok=True) moved = 0 for fname, label in APPLE_LABELS.items(): src_path = os.path.join(src_dir, fname) dst_path = os.path.join(dst_dir, label, fname) if os.path.exists(src_path): shutil.copy2(src_path, dst_path) moved += 1 else: print(f"WARNING: {src_path} not found") print(f"moved {moved} files into {dst_dir}")这段代码里有一个关键决定:用shutil.copy2而不是shutil.move。原因是我在上一节说过要保持原始文件不动,复制出来的是整理后的副本,这样万一目录组织方式出错,原始文件还能救回来。copy2会比copy多复制元数据(包括时间戳),我一般习惯用它,方便后续比对原始文件的修改时间。
目录结构就变成了这样:
apple_dataset/ ├── fuji/ │ ├── 203.jpeg │ ├── 18.jpeg │ └── ... ├── gala/ ├── granny_smith/ └── red_delicious/这个结构下,ImageFolder会自动按目录名的字母序分配 class index,比如fuji是 0,gala是 1。这个索引顺序在后续评估时要特别注意——打印混淆矩阵时行列名称很容易和这批索引对不上。
3.3 划分 train / val / test 并生成清单
目录结构就位后,紧接着做数据划分。划分比例我会用 8:1:1,即 464 张训练、58 张验证、58 张测试。这个比例对 580 张的数据集来说偏保守,但比 9:1(522 张训练、58 张测试)更稳,因为验证集每个月都在调参时被反复看,验证集和测试集必须彻底分家,否则测试集就不再是「没见过的数据」了。
import os import random import shutil random.seed(42) src_dir = "apple_dataset" train_dir = "split_dataset/train" val_dir = "split_dataset/val" test_dir = "split_dataset/test" categories = os.listdir(src_dir) for split_dir in (train_dir, val_dir, test_dir): for c in categories: os.makedirs(os.path.join(split_dir, c), exist_ok=True) for c in categories: files = os.listdir(os.path.join(src_dir, c)) random.shuffle(files) n = len(files) n_train = int(n * 0.8) n_val = int(n * 0.9) # 剩下的 10% 归测试 for i, f in enumerate(files): if i < n_train: shutil.copy2(os.path.join(src_dir, c, f), os.path.join(train_dir, c, f)) elif i < n_val: shutil.copy2(os.path.join(src_dir, c, f), os.path.join(val_dir, c, f)) else: shutil.copy2(os.path.join(src_dir, c, f), os.path.join(test_dir, c, f)) print("done")random.seed(42)的作用是让划分结果可复现——同一份数据、同一个 seed,每次运行得到的划分边界完全一致,这样不同人调参时对比的基线才是公平的。random.shuffle先打乱再切分,避免某个品种的图片在原始文件系统里恰好按拍摄时间排序,导致前面 80% 全是同一批次拍摄的同角度图。
这里有一个必须注意的细节:数据划分必须在「品种」维度上分层,而不是在所有 580 张图上打乱后切片。上面代码是按类别目录逐个处理的,每个品种内部独立划分,这保证了任何一个品种的图片都同时出现在 train、val、test 三个集合中。如果全局打乱再切,某些小样本品种可能整批掉进训练集,测试集里一个都没有——模型对这个品种的识别准确率直接变 0。
4. 基于 PyTorch 训一个苹果品种分类模型:ResNet18 微调全程
4.1 为什么选 ResNet18 而不是自己搭网络
数据规模只有 580 张,这是一个非常关键的信息。580 张图对深度学习来说属于典型的小规模数据集,从头训练一个 VGG 或者自研 CNN 很容易过拟合——模型参数量远远大于样本量,训练集准确率能冲到接近 100%,验证集却一直卡在 70% 上下晃。这个场景下,迁移学习是更稳的选择,用 ImageNet 上预训练过的 ResNet18 做骨干,冻结前几层,只微调最后几层和全连接头,既能吃预训练模型已经学到的边缘、纹理、形状特征,又不会因为数据量不够而把底层特征学偏。
ResNet18 相对 ResNet50 的优势在于参数量小(约 11M 参数),在只有几百张图、大概率没有顶级 GPU 的环境下训练更快,而且残差连接的结构对梯度的传播更友好,不容易出现浅层梯度消失。如果你的环境里有 16GB 以上显存,换 ResNet50 也可以,但对这个数据规模,ResNet18 的效果和 ResNet50 差距很小,训练时间却少一半。
4.2 数据增强参数与加载器实现
数据增强在小数据集上的作用比在大数据集上更明显。580 张训练图,如果不做增强,模型每轮看到的只有这 464 张,几个 epoch 下来就背下来了;做了随机裁剪、翻转、颜色抖动之后,模型每轮看到的都是「同一张图的不同状态」,相当于把数据量放大了好几倍。下面是这套流程里我会用的增强配置:
from torchvision import transforms train_transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) val_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]) ])参数逐个说。Resize((256, 256))先把图放大到 256,RandomResizedCrop(224, scale=(0.7, 1.0))再从 256 的图里随机裁一块 224×224 的区域,scale=(0.7, 1.0)表示裁剪面积占原图的 70% 到 100%,这样模型能看到苹果在不同尺度下的形态,对拍摄距离的变化更鲁棒。RandomHorizontalFlip(p=0.5)以 50% 概率水平翻转,苹果是左右对称的物体,水平翻转不会破坏语义,但会让模型学习到更对称的特征。RandomRotation(15)旋转角度控制在正负 15 度,苹果数据集里的图片大多是平放拍摄,旋转超过 15 度会让苹果看起来像是「要滚下桌子」,反而增加学习难度。ColorJitter的四个参数分别控制亮度、对比度、饱和度和色相的随机扰动幅度,hue 扰动只设 0.05,因为苹果的颜色(红、绿、黄)本身就是品种区分的关键特征,色相扰动太大会让红苹果变成黄苹果,直接破坏标签语义。
验证集和测试集只用Resize((224, 224))加归一化,不做任何随机增强,保证评估结果稳定可复现。mean和std用的是 ImageNet 统计值,因为预训练模型是在 ImageNet 上训练的,输入数据用 ImageNet 的归一化参数最合适。
4.3 训练代码主体与参数调节
训练脚本是这套流程的核心,我把完整可跑的版本贴出来:
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, models batch_size = 16 num_epochs = 30 learning_rate = 1e-4 num_classes = len(os.listdir("split_dataset/train")) train_dataset = datasets.ImageFolder("split_dataset/train", transform=train_transform) val_dataset = datasets.ImageFolder("split_dataset/val", transform=val_transform) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) model.fc = nn.Linear(model.fc.in_features, num_classes) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=3)这里有几个参数值得展开。batch_size = 16对 580 张的数据集来说是合适的值,ResNet18 在 224×224 输入下,单卡 8GB 显存跑 batch 16 没问题;如果你显存紧张可以降到 8,但学习率也要相应调小。learning_rate = 1e-4是微调预训练模型时的常用起点,比从头训练用的 1e-3 要低一档,因为预训练特征已经接近最优解,学习率太大会破坏已经学好的底层特征。ReduceLROnPlateau是学习率调度器的关键选择:当验证集 loss 连续 3 个 epoch(patience=3)不下降时,学习率乘以factor=0.5减半,这比固定步长衰减更适配实际训练曲线——前期 loss 快速下降,后期曲线进入平台期,学习率自动缩小,帮助 loss 继续探底。
训练循环本体:
best_val_acc = 0.0 for epoch in range(num_epochs): model.train() running_loss = 0.0 correct = 0 total = 0 for inputs, labels in 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() * inputs.size(0) _, preds = torch.max(outputs, 1) total += labels.size(0) correct += (preds == labels).sum().item() train_loss = running_loss / total train_acc = correct / total val_loss, val_acc = evaluate(model, val_loader, criterion, device) scheduler.step(val_loss) if val_acc > best_val_acc: best_val_acc = val_acc torch.save(model.state_dict(), "best_apple_model.pth") print(f"Epoch {epoch+1}/{num_epochs} | " f"train_loss {train_loss:.4f} train_acc {train_acc:.4f} | " f"val_loss {val_loss:.4f} val_acc {val_acc:.4f}")训练循环的逻辑很直接:每个 epoch 分训练和验证两个阶段,训练阶段开着 dropout 和增强,验证阶段切到model.eval()关闭这些随机性。optimizer.zero_grad()每批次清零梯度,避免上一个 batch 的梯度累积到当前 batch。保存模型用的是model.state_dict()而不是整个模型对象,后者会把模型结构、优化器状态全打包,文件大且换环境时容易版本不兼容;state_dict()只存参数,配合加载侧重新定义模型结构,跨环境恢复权重方便得多。
evaluate函数需要自己补全,逻辑很固定:
def evaluate(model, loader, criterion, device): model.eval() running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for inputs, labels in loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) running_loss += loss.item() * inputs.size(0) _, preds = torch.max(outputs, 1) total += labels.size(0) correct += (preds == labels).sum().item() return running_loss / total, correct / total注意验证阶段必须包在torch.no_grad()里。no_grad会关闭自动求导的图构建,推理速度提升明显,而且省显存——不关的话,每个 batch 的前向计算都会把中间结果缓存下来供反向传播使用,验证阶段根本不需要反向传播,这些缓存纯属浪费。
5. 苹果分类数据集落地避坑:损坏图、脏标签与过拟合
5.1 解压或读取时遇到「损坏图片」报错
现象:用ImageFolder加载数据集时,训练跑到一半突然报错,提示类似OSError: cannot identify image file或PIL.UnidentifiedImageError,有的直接崩进程。
原因:这份数据集以 zip 包形式传播,压缩包里部分 JPEG 文件可能在压缩或传输过程中损坏。Image.open()是惰性加载,文件头没问题就能打开,但DataLoader里的default_loader会完整读取像素数据,损坏的区域一访问就抛异常。
解决:先跑第 2.2 节里的预审脚本,把损坏图片全部找出来。对被标记的图片,先看有没有备用压缩包,重新解压覆盖;如果原始包里就缺,那就把该图片从映射表和目录里剔除,并在训练前手动把类别样本数重新统计一遍。如果是压缩包带密码或嵌套压缩导致解压不全,先验证 MD5 确认包完整性,不要把希望寄托在各种所谓的「zip 密码移除」工具上,那些工具对正儿八经加密的 zip 基本无效。
5.2 训练 loss 不降,或验证集准确率极低
现象:训练了 10 个 epoch,训练集 loss 一直在 2.3 左右徘徊(接近随机猜测的 log(类别数)),验证集准确率始终低于 40%;或者训练集 loss 降得很快,验证集 loss 反而上涨。
原因:训练 loss 不降,最常见的原因是标签和图片错位——比如映射文件里把203.jpeg标成了fuji,但 203 号图片实际是青苹果,模型被喂了矛盾的数据,学不到有效特征。另一种常见原因是从ImageFolder读取的类别索引和标签名称对不上,ImageFolder按目录名字母序分配索引,fuji是 0、gala是 1,但你在评估时拿品种名字典去查,索引一错位,准确率计算就是错的。
解决:训练前用第 3.2 节的校验脚本,确保映射表里的所有 key 都在磁盘上存在,所有磁盘上的文件都在映射表里有标签。训练到第 3 个 epoch 时,随机抽一个 batch 的输入图片和标签打印出来,人工核对——左边是图,右边是类别名,肉眼扫一眼就能发现错位。索引对齐问题,在模型保存后先跑一次推理,拿几个已知品种的测试图片,把输出索引映射回品种名,确认映射关系正确后再看准确率。
5.3 小样本品种过拟合严重,验证集震荡
现象:训练曲线震荡明显,某个品种的训练准确率是 100%,验证集准确率只有 50%;换一个随机种子,结果上下波动超过 10 个百分点。
原因:580 张图平均分到多个品种后,某些品种可能只有 30 张甚至更少,训练集里只有 24 张。这么小的样本量,模型很容易把这 24 张图的背景、拍摄角度、光线条件全部记住,而不是真正学会苹果品种的形态特征。验证集波动大,恰恰说明划分时没有做分层,某个品种的验证图片恰好都是难的样本,或者验证图片数量太少,随机性主导了评估结果。
解决:首先确认第 3.3 节的分层划分做对了,每个品种独立切分。其次给损失函数加类别权重,或者用WeightedRandomSampler在 DataLoader 层做上采样。WeightedRandomSampler的做法是计算每个类别的样本数倒数作为采样权重,让模型每个 batch 里小样本类别出现的概率更高:
from torch.utils.data import WeightedRandomSampler labels = [train_dataset.targets[i] for i in range(len(train_dataset))] class_counts = torch.bincount(torch.tensor(labels)) weights = 1.0 / class_counts.float() sample_weights = weights[labels] sampler = WeightedRandomSampler(sample_weights, num_samples=len(train_dataset), replacement=True) train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler)replacement=True表示同一个样本可以在一个 epoch 里被重复采样,这样小样本类别在一个 epoch 里的出现次数能被拉到大样本类别的量级。代价是训练一个 epoch 的时间变长,但对 580 张图来说完全可接受。
5.4 图片方向错乱导致识别率上不去
现象:模型对所有品种的验证集准确率都卡在 85% 左右上不去,看起来像是欠拟合,但加大训练轮数也没用。把几张预测错误的图片打印出来发现,苹果是倒着放的或者旋转了 90 度,肉眼还能认出来,模型却认错了。
原因:手机和相机拍摄时,EXIF 里记录了方向信息(orientation tag)。部分图片在电脑上预览时显示正常,是因为看图软件自动读取了 EXIF 的方向并做了旋转;但 PyTorch 的default_loader用的是PIL.Image.open(),它不会自动处理 EXIF 方向,导致一部分图片在训练时是旋转过的状态。
解决:在 Dataset 的__getitem__里用ImageOps.exif_transpose修正方向:
from PIL import ImageOps def load_image(path): with Image.open(path) as im: im = ImageOps.exif_transpose(im) return im.convert("RGB")exif_transpose会读取 EXIF orientation 字段,把图片物理旋转到正确方向后再返回。这一步必须在 Resize 和增强之前做,顺序错了等于没做。如果你的数据集是从网上爬的或者经过二次压缩,EXIF 可能已经被去除,这一步不会产生任何变化,但做了能消除一个隐藏变量。
6. 验证阶段别只看准确率:混淆矩阵与误判样本回看
准确率这个指标,在多品种分类任务里会骗人。假设这份数据集里富士苹果的样本占了一半,模型把所有图都预测成富士,准确率也能有 50%,但这个模型在生产环境里毫无价值。我一般会在训练完成后做两件事:画混淆矩阵,然后把误判样本打印出来逐张看。
import torch import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import confusion_matrix from torchvision import datasets, models model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for inputs, labels in test_loader: inputs = inputs.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm = confusion_matrix(all_labels, all_preds) class_names = test_dataset.classes plt.figure(figsize=(8, 6)) sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=class_names, yticklabels=class_names) plt.xlabel("Predicted") plt.ylabel("True") plt.tight_layout() plt.savefig("confusion_matrix.png", dpi=150)fmt="d"表示矩阵里的数字按整数显示而不是科学计数法,小样本数据集里一张图的差异都可能在混淆矩阵里表现为一个很刺眼的非零格子。xticklabels和yticklabels用test_dataset.classes,确保每个格子对应的类别名称是正确的,避免索引错位。混淆矩阵里对角线越亮越好,非对角线上的亮点就是品种间最容易混淆的组合——比如把gala认成fuji,这两类都是红色系苹果,外形本来就接近。
看完矩阵再看误判样本的具体图片,这一步能帮你判断是数据问题还是模型问题。我通常会打印每个误判的样本,把预测类别、真实类别和图片一起输出:
import numpy as np from PIL import Image misclassified = np.argwhere((np.array(all_preds) != np.array(all_labels))).flatten() for idx in misclassified[:10]: true_label = class_names[all_labels[idx]] pred_label = class_names[all_preds[idx]] img_path = test_dataset.samples[idx][0] img = Image.open(img_path).resize((224, 224)) img.save(f"misclass_{idx}_{true_label}_as_{pred_label}.jpg") print(f"{img_path} | true: {true_label} | pred: {pred_label}")误判样本存在本地后,逐张打开看,重点观察两个细节。第一,图片里苹果是否完整,如果有遮挡、切边严重、或者一个画面里有多个苹果,模型判断出错可以归因于数据质量,而不是模型能力;第二,误判组合是否集中在颜色相近的品种之间,如果是,说明颜色特征在模型决策里权重过高,需要考虑增加形状或纹理相关的预处理,比如把图片转成灰度通道叠加到 RGB 上再训练。
这套验证流程看起来繁琐,但每次执行都能发现新问题。我从那以后每次拿到类似的数据集,都会在训练完成后强制走一遍「预审 → 分层划分 → 分类权重 → 混淆矩阵 → 误判回看」的完整流水线,哪怕数据只有 100 张也不跳过。数据集的坑永远不会出现在报告里,只会在训练曲线和混淆矩阵里露出马脚。希望帮到你。
本文还有配套的精品资源,点击获取