ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch实战:36种果蔬分类识别的CNN模型与视频推理优化

PyTorch实战:36种果蔬分类识别的CNN模型与视频推理优化

1. 一上来先聊聊这个项目是干什么的

如果你经常下厨或者管理生鲜供应链,应该能理解这个痛点:一堆水果蔬菜堆在那里,人工分类费时费力,而且不同品种长得又像,比如苹果和梨、西兰花和花菜,靠肉眼分辨很容易出错。所以用深度学习做一个自动识别的系统,在实际场景里非常有用。

这个项目的核心是构建一个基于PyTorch框架的深度卷积神经网络模型,能够对36种常见的水果和蔬菜进行分类识别。模型需要支持三种输入形态:单张照片识别、照片批量连续识别、实时视频流识别。整体难度属于计算机视觉中的经典多分类任务,上手门槛不高,但要做到识别的稳定性和实时性,里面还是有不少门道。

这个项目适合三类人参考:一是刚学完PyTorch基础、想找一个完整CV项目练手的人;二是做智慧农业、生鲜零售、餐饮结算系统相关开发的技术人员;三是想了解如何把训练好的图像分类模型落到视频实时场景的工程师。我会把整个实现链路——从数据预处理到模型训练,再到视频推理加速——完整拆开讲清楚,最后附上我实际踩过的坑和排查方法。

2. 整体方案设计:为什么选PyTorch、为什么用深度卷积网络

2.1 框架选型的真实对比

很多初学者会纠结选PyTorch还是TensorFlow,我直接说我的个人结论:在这个项目上,PyTorch更顺手。

原因有三点。第一,PyTorch的动态计算图让调试变得直观。训练过程中你随时可以打印中间张量的维度、检查梯度是否正常,这对于排查模型不收敛的问题帮助极大。第二,PyTorch生态里的torchvision自带大量预训练模型(ResNet、DenseNet、EfficientNet等),而且权重下载稳定,迁移学习非常方便。第三,社区里关于PyTorch的CV教程数量和质量都占优,遇到问题更容易搜到答案。

当然,如果你所在团队已经有成熟的TensorFlow部署体系,那用TensorFlow也完全可以,因为这个任务本身并不依赖特定框架。但单从开发效率和调试体验来说,PyTorch对于动手实践类的项目是明显更合适的。

2.2 为什么任务能落地:36分类的难度评估

36种果蔬的分类,在深度学习里属于中等规模分类任务,不算难,但也别小看了。

难点主要在类间相似和类内差异。比如不同品种的苹果,颜色、形状差异很大,而青苹果和梨站在一起,轮廓和颜色又很像。再比如柠檬与青柠。这些情况要求模型不能只依赖颜色和形状的浅层特征,而是要捕捉到更细微的纹理和结构差异,这就对网络容量和数据质量提出了要求。

用深度卷积网络来应对这个任务是比较合适的。卷积神经网络天然具备局部感知和权值共享的特性,能够通过层层叠卷积核提取从边缘、纹理到形状、部件的层次化特征。相比早期手工设计特征(颜色直方图、HOG、SIFT)的做法,CNN免去了大量特征工程工作,而且泛化能力明显更好。

2.3 方案架构的总体设计

整个系统分为两个阶段:离线训练阶段和实时推理阶段。

离线训练阶段,我们用预训练的卷积神经网络模型作为主干,替换最后的全连接分类层,用果蔬数据集进行迁移学习。训练完成后导出模型权重文件。

实时推理阶段,分为图片模式和视频模式。图片模式是单张读入、预处理、模型推理、输出结果。视频模式本质上就是图片模式的循环执行,但需要考虑帧率控制和性能优化。为了让视频识别跑得流畅,推理阶段通常会做输入尺寸缩放、半精度推理(FP16)、以及每隔N帧采样一次的策略。

整体流程图不画了,文字描述就是:视频帧 → 帧采样 → 预处理(缩放/归一化/转张量) → 模型前向传播 → softmax得到概率分布 → 取top-1或top-3类别 → 在帧上绘制标签 → 输出视频。这条链路在代码实现上并不复杂,但每一步都有对应的工程细节需要处理。

3. 数据集准备与预处理:36种果蔬的数据从哪来、怎么生成

3.1 数据集构建思路

做36分类,每类至少需要几百张有效图片,否则很容易过拟合。这里有两种路径:

路径一,使用公开数据集。Kaggle和GitHub上有不少果蔬分类数据集,有的包含几十种水果蔬菜,每类上百到上千张图。比如常见的有Fruits 360这类数据集,类别数量足够覆盖到30多种。使用公开数据集的优点是省事,缺点是某些类别可能和你实际要识别的品种不完全一致,需要自己做筛选和清洗。

路径二,自建数据集。用手机或相机拍摄食材,注意拍摄角度、光照、背景要有变化。实际做的时候我建议把两种方式结合:以公开数据集为主,再自己补充一些在真实场景中会遇到的图片(比如超市包装盒里的、带泥土的、切片后的),这样能明显提高模型的实战效果。

3.2 数据清洗的关键操作

数据清洗是最容易被低估的环节。几个常见问题:

  • 重复图片:有的公开数据集里存在一模一样的文件,直接训练会让模型对重复样本过拟合,要用哈希去重(计算图片的md5或感知哈希)清理。
  • 错误标签:数据集里偶尔会有标签错乱的情况,比如把青柠标成柠檬。如果训练集里混入这类错误样本,模型会学错。建议每类图片抽样人工看一遍,至少要抽查。
  • 模糊与极端图片:完全失焦、严重过曝的图片对训练没有帮助,直接剔除。

3.3 数据增强策略与具体参数

为了让模型在真实场景中更鲁棒,仅靠原始图片是不够的。PyTorch的torchvision.transforms提供了丰富的数据增强工具,我在项目中用的核心增强组合如下:

增强操作参数设置作用
RandomResizedCrop缩放范围0.7-1.0,尺寸224x224模拟不同距离拍摄的尺度变化
RandomHorizontalFlip概率0.5模拟左右角度变化
RandomRotation角度±15度模拟摆放角度差异
ColorJitterbrightness=0.2, contrast=0.2, saturation=0.2模拟光照和色彩环境变化
Normalizemean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]与预训练模型输入要求对齐

这里特别注意Normalize的数值不能乱改。因为使用预训练模型时,主干网络是在ImageNet上训练的,它的输入数据分布就是上面这套均值和标准差。如果归一化方式和预训练不一致,特征分布会偏移,效果大打折扣。

3.4 数据划分与目录结构

数据按train/val两目录划分,每类一个子文件夹,这是PyTorch的torchvision.datasets.ImageFolder能够直接读取的标准目录结构:

data/ ├── train/ │ ├── apple/ │ │ ├── apple_001.jpg │ │ ├── apple_002.jpg │ │ └── ... │ ├── banana/ └── val/ ├── apple/ ├── banana/ └── ...

划分比例使用8:2,训练集和验证集需要注意在类别维度上保持均衡。如果某个原始类别图片数量特别少,优先回到上一步去补充数据,而不是强行训练,否则验证集指标会虚高。

4. 模型构建:迁移学习的准确选型与实践

4.1 主干网络挑选

36种果蔬分类,输入图像的分辨率不算高,原始图片基本是几百像素的级数,所以主干网络不追求极致规模,而要平衡精度和推理速度。

我当时对比了几种常见backbone的实际表现:

网络模型参数量级推理速度(CPU)Top-1准确率表现适用场景
ResNet18较小较快中等必要时的轻量级选择
ResNet50中等中速较高均衡的推荐选择
EfficientNet-B0中等中速较高精度与速度都不错
DenseNet121中等中速较高特征复用能力强

我的最终选择是ResNet50。原因有几个:第一,ResNet的残差结构在训练稳定性方面表现很好,不容易梯度消失;第二,torchvision里的resnet50提供了预训练权重,迁移学习起步精度高;第三,它的推理速度在GPU上非常快,在普通CPU上也能跑,兼顾了实时视频场景。

如果你部署的机器性能较弱(比如树莓派或老笔记本),可以换ResNet18或者MobileNetV3,精度会低几个点,但帧率能上去;相反,如果设备性能强劲且对精度有更高要求,可以尝试EfficientNet-B3或更大版本。

4.2 迁移学习核心操作

使用torchvision加载预训练ResNet50后,关键一步是替换最后的全连接层:

import torchVision.models as models model = models.resnet50(pretrained=True) num_features = model.fc.in_features # 36类分类,最后一层替换 model.fc = torch.nn.Linear(num_features, 36)

这里不要直接新建一个nn.Linear(2048, 36)就完事了,一定要先通过model.fc.in_features取出原来预训练模型的输入维度(ResNet50是2048),避免硬编码写错。这个细节在很多参考代码里容易被忽略,但对齐不一致会导致forward时报维度错误。

4.3 冻结策略与解冻时机

迁移学习中一个常见策略是:先冻结主干网络的所有参数,只训练分类头(全连接层)。这样训练快,因为backbone的预训练特征已经很强了,只需要让分类头去学习映射。

但如果你只训练分类头,精度会很快饱和。我建议的流程是:

阶段一:冻结backbone,lr设置为0.001,训练5个epoch,让分类头稳定下来。阶段二:解冻一部分backbone(通常是最后几层block),lr降到0.0001,继续训练10到15个epoch,微调高层特征。阶段三:用更低的学习率(0.00001)和更小的数据增强幅度(去掉RandomRotation,保留ColorJitter),再训练几个epoch,让模型在验证集上收敛。

注意,解冻层数不是越多越好。如果全部解冻且lr用错,预训练权重会被粗暴冲掉,效果可能还不如直接用随机初始化从头训练。

4.4 损失函数与优化器选择

36分类属于多分类单标签问题,标准搭配是交叉熵损失(CrossEntropyLoss)。

优化器我从实际经验出发推荐两档:

  • Adam:训练初期收敛快,调试方便,适合快速验证模型是否正常工作。
  • SGD(momentum=0.9):到后期微调阶段,SGD往往能取得比Adam更低的验证损失,但需要手动调学习率,训练节奏更需要耐心。

我的做法是:前两个阶段用Adam,最后解冻微调阶段切换到SGD。切换时学习率要降一个量级(比如从1e-4降到1e-5),并且不要沿用Adam的动量状态,直接新建优化器。这个切换逻辑在工程上虽然有点"野",但实测效果通常优于全程只用一种优化器。

5. 训练过程与调优细节:损失曲线的识别和超参调整

5.1 训练循环的核心代码框架

训练循环本身并不复杂,关键是细节要到位。核心结构如下:

for epoch in range(num_epochs): model.train() running_loss = 0.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) # 每个epoch跑一次验证 val_loss, val_acc = validate(model, val_loader, criterion, device) print(f'Epoch {epoch+1}: train_loss={running_loss/len(train_dataset):.4f}, ' f'val_loss={val_loss:.4f}, val_acc={val_acc:.4f}')

训练中不要忘了在验证阶段切换model.eval(),并且用with torch.no_grad():包裹前向推理,否则PyTorch会默认跟踪梯度,白白增加内存和显存消耗。这个坑几乎每个新手都会踩,最明显的表现就是验证集跑得很慢、显存爆掉。

5.2 学习率调度策略

学习率调度建议使用ReduceLROnPlateau,它的策略是:当验证损失连续若干个epoch不再下降时,学习率自动乘以一个因子(如0.1)。

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.1, patience=3, verbose=True ) # 每个epoch验证完之后 scheduler.step(val_loss)

patience=3表示连续3个epoch验证损失没有刷新最低值,就降低学习率。verbose=True会打印学习率变化信息,方便观察。这个策略好过固定步长的StepLR,因为它能自适应任务的收敛节奏,不用人工盯着曲线去判断何时降学习率。

5.3 让训练更稳定的两个小技巧

技巧一:标签平滑。把交叉熵目标的硬标签(独热编码)替换为软标签,例如对正确类别给0.95,其余给0.05除以35。这样可以防止模型对训练集过度自信,减轻过拟合。实现方式是在CrossEntropyLoss中设置label_smoothing=0.1(PyTorch 1.10以上版本支持)。

技巧二:早期停止。保存验证准确率最高的模型权重(best_model.pth),而不是最后一个epoch的权重。一个简单的做法是每次验证结束后判断val_acc是否高于历史最优,是则保存:

if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), 'best_model.pth')

5.4 训练过程中的常见误区

观察训练曲线时,有几个典型的"假象"要识别:

验证集正常、训练集损失也在降,但召回率在部分类别上特别低。这是类别不平衡的典型信号,说明个别类别的数据量太少或特征太接近。处理方法优先考虑补充数据,而不是修改损失函数权重。

验证损失上下震荡厉害。通常是学习率太大,或者batch size太小导致梯度噪声过大。先尝试降低学习率一个量级,并把batch size提高到32或64。

训练集准确率接近100%,验证集却在80%附近徘徊。这是标准的过拟合。检查增强策略是否足够、是否解冻了过多层、是否训练了太多epoch。通常的做法是增加增强强度、增加dropout、减少解冻层数。

6. 照片连续识别与视频识别的实现

6.1 单张照片推理流程

模型训练好后,推理阶段我们写一个独立的脚本。

def predict_image(img_path, model, transform, class_names): image = Image.open(img_path).convert('RGB') image = transform(image).unsqueeze(0) image = image.to(device) with torch.no_grad(): outputs = model(image) probs = torch.nn.functional.softmax(outputs, dim=1) top3 = probs.topk(3, dim=1) result = [(class_names[idx], float(score)) for idx, score in zip(top3.indices[0], top3.values[0])] return result

推理时需要注意将图片转换为RGB三通道。很多手机照片使用EXIF旋转信息,直接读取时可能是旋转前的状态,建议在预处理阶段把Image.open(...).convert('RGB')之后,额外做一次EXIF方向的纠正(可以通过PIL的ImageOps.exif_transpose实现)。这个问题不处理的话,偶尔会出现同一张竖拍照片在多个角度下识别结果不一致的情况。

6.2 照片批量连续识别

连续识别本质上就是遍历文件夹中所有图片,逐张推理并汇总结果。这里有两层优化思路。

第一层是批量推理而非循环单张推理。把一批图片同时送入GPU,利用矩阵运算的并行能力提升吞吐量。

def predict_batch(image_paths, model, transform, class_names, batch_size=32): results = [] for i in range(0, len(image_paths), batch_size): batch_paths = image_paths[i:i+batch_size] batch_tensors = [] for p in batch_paths: img = Image.open(p).convert('RGB') batch_tensors.append(transform(img)) batch_tensor = torch.stack(batch_tensors).to(device) with torch.no_grad(): outputs = model(batch_tensor) probs = torch.nn.functional.softmax(outputs, dim=1) top_idx = probs.argmax(dim=1) for p, idx, prob in zip(batch_paths, top_idx, probs.max(dim=1).values): results.append((p, class_names[idx.item()], prob.item())) return results

第二层是并发处理。如果图片数量极大,而且每张图需要做很重的预处理,多进程读取图片、主进程负责模型推理是一个常见模式。对于入门项目,我建议先用单进程的实现,如果吞吐量不够再改并发。

6.3 视频识别:从图片到视频流的转换

视频识别的核心机制是用OpenCV读取摄像头或视频文件,逐帧提取,并把帧转换为模型需要的张量格式。

import cv2 import torch from PIL import Image cap = cv2.VideoCapture(0 if use_camera else video_path) frame_interval = 3 # 每隔3帧处理一次 frame_count = 0 while cap.isOpened(): ret, frame = cap.read() if not ret: break frame_count += 1 if frame_count % frame_interval != 0: continue # OpenCV的BGR图转换为RGB PIL Image frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) pil_img = Image.fromarray(frame_rgb) # 预处理与推理 input_tensor = transform(pil_img).unsqueeze(0).to(device) with torch.no_grad(): outputs = model(input_tensor) probs = torch.nn.functional.softmax(outputs, dim=1) top1_idx = probs.argmax(dim=1).item() confidence = probs.max(dim=1).values.item() # 绘制结果并显示 label = f'{class_names[top1_idx]} {confidence:.2f}' cv2.putText(frame, label, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) cv2.imshow('Fruit and Vegetable Recognition', frame) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows()

6.4 视频识别性能优化三板斧

视频识别最容易出现的问题是画面卡顿,因为推理速度跟不上视频帧率。优化方案有两个方向:

方向一:降低单帧推理耗时。

  • 输入尺寸压缩:视频帧可以从1280x720或1920x1080直接缩放到256x256或224x224。因为这个任务本身不需要超高清细节,缩小时信息损失不大,推理时间能减少一半以上。
  • 半精度推理:在支持FP16的GPU上,将模型转为半精度:
model.half()

注意输入张量也要跟着转成half类型。

  • TorchScript或ONNX导出:对CPU部署尤其有效,ONNX Runtime推理比PyTorch原生CPU推理通常快20%到30%。

方向二:降低推理频率。

食品识别场景中,果蔬不会快速运动,完全没必要每帧都推理。按间隔N帧或间隔固定时间采样即可。上面代码里frame_interval=3意味着每3帧只处理1帧,显示上依然是接近实时的,但模型负载直接降到1/3。如果识别结果变化不频繁,还可以增加一个"结果稳定"逻辑:只有当连续多次采样得到高置信度类别时才更新画面上的标签,避免标签闪烁。

方向一和方向二可以叠加使用。在NVIDIA GPU上,224x224输入、ResNet50、FP16推理,单帧耗时大约在5到10毫秒,加上预处理和绘制,基本可以做到30FPS的视频流实时识别。CPU环境下,按5帧间隔采样,也能跑到接近可用的效果。

6.5 视频识别中帧预处理的一个隐藏坑

OpenCV读出来的帧是BGR色彩空间,而torchvision预训练模型期望的是RGB。很多人在代码里直接用cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)做了转换,但遗漏了另一件事:归一化。

transforms里如果包含了Normalize,那在把PIL图像转为Tensor之后,Normalize会做减均值除标准差的运算,这一步在标准transform流程中是自动完成的。但你如果用自定义的方式(比如直接对numpy数组做归一化),就必须手动对齐。最稳妥的写法是统一走transform流程,不要自己写预处理逻辑,既简化代码又不容易出错。

7. 常见问题与排查技巧实录

7.1 视频推理帧率太低怎么排查

先分清瓶颈在哪一步。给代码加上计时分析和日志,分别统计读帧耗时、预处理耗时、模型推理耗时、绘制耗时。根据我的经验,在CPU上,模型推理通常占大头;在GPU上,读帧和预处理可能成为瓶颈。

如果瓶颈在推理,优先降低输入尺寸、换更轻量级的backbone、转为ONNX Runtime。如果瓶颈在读帧或预处理,使用cv2.VideoCapture时设置CAP_PROP_BUFFERSIZE优化缓冲区,或者把预处理放到单独的线程去执行。

7.2 训练时验证集准确率上不去

分情况排查:

  • 检查数据增强是否在验证集生效:验证集的transform中不应包含RandomResizedCrop、RandomHorizontalFlip等随机操作。初学者容易直接在验证集沿用训练集的transform,导致验证集结果不稳定。
  • 检查标签与类别是否对齐:torchvision.datasets.ImageFolder的类别顺序是字典序。你在推理时构建class_names列表必须与这个顺序一致。如果顺序错乱,验证集准确率看着还行、但实际推理结果对不上,就是因为两者类别映射不一致。
  • 检查是否加载了最优权重文件:训练完成之后,经常有人加载的是last_model.pth而不是best_model.pth。如果训练后期出现过拟合,最后一个epoch的权重明显不如历史最优权重。

7.3 显存不足(CUDA Out of Memory)

最直接的解决办法是我前面提到的批量大小调小、图片尺寸调小。另外注意训练循环里是否有梯度累积未清零的情况。每个batch结束后一定要optimizer.zero_grad(),否则梯度会累积到下个step,既影响收敛也会占用额外内存。

视频推理时如果遇到显存不足,多半是因为用了一个长期的VideoCapture循环但历史帧依然被张量引用。可以用del frame, input_tensor手动释放,或者把推理封装成函数,避免局部变量累积。

7.4 一些类别的准确率明显更低

比如在实验中苹果类别整体识别不错,但"蛇果"类会被误判成"苹果";"西红柿"又会被误判为"圣女果"。处理这种问题我的经验是:优先做细粒度数据的扩充,多收集易混淆类别之间差异明显的图片,比如侧面、切面、特写。其次,在类别名和训练数据的组织上尽量体现细粒度差异,让分类器有更多判别依据。

如果数据实在没法扩充,也可以在类别名称上做合并,比如把容易混淆的品种归为同一大类。工程上始终要记得:模型是服务于落地场景的,45种和36种并不重要,识别用户真正关心的类别并给出可靠结果才是关键。

8. 部署扩展与个人经验分享

8.1 USB摄像头与网络摄像头的接入方式

如果想在树莓派或开发板上跑模型,摄像头读取这一块要格外注意。USB摄像头通常用OpenCV的cv2.VideoCapture(0)即可驱动,而RTSP协议的网络摄像头(IP Camera)则需要设置cv2.CAP_FFMPEG或cv2.CAP_ANY,并传入RTSP地址。RTSP流在不稳定的网络环境下容易出现断帧,建议在读取循环中加入重连机制。

8.2 从脚本到服务化的思路

如果这个识别功能要给别的系统调用,最简单的做法是封装一个HTTP接口。PyTorch模型被Flask或FastAPI加载一次后,在每个请求中执行推理。请求体里传图片文件或base64编码,响应中返回类别、置信度和耗时。需要强调的是,模型加载应该放在服务启动阶段,不要在每次请求中重复加载,否则内存和延迟都会变差。

8.3 我完整跑通这个项目后的真实收获

这个项目我完完整整跑通过一遍,最大的感受就是:一个看似简单的分类任务,在"单张图片"和"实时视频"之间存在不小的工程鸿沟。单张图片识别做得好,不代表视频识别效果好。视频引入了一个新维度:时间。帧率、延迟、抖动、光照变化,都会让你的模型暴露出在单张测试时看不到的问题。

我个人在整个过程中学到最重要的一件事是:性能优化永远是面向场景的。在一个果蔬识别的场景里,用户大概率不会在一个画面停留不到一秒钟就要求模型出结果。所以帧采样、结果平滑、置信度阈值这套"实时工程的组合拳",比单纯追求模型精度更能让整个系统变得可用。

8.4 后续可以这样扩展

这个项目上有几个值得做的扩展方向:

  • 加入检测功能:现在做的是全局分类,即整个画面被认为只有一种果蔬。真实场景中画面里可能同时有多种果蔬,这时需要在此框架上叠加目标检测器(YOLO或Detectron2),先定位每个目标的位置,再独立分类。
  • 加入类别拒识:当前模型会对任意输入返回36类之一。如果画面里出现一个杯子,模型会强行归到果蔬类别之一,使用软最大概率的置信度作为门槛可以在一定程度上过滤此类错误。
  • 加入类别增量学习:36类跑通后想扩展到50类,不需要从头训练,可以在现有模型基础上微调新增类别,但要注意旧类别的灾难性遗忘问题,需要保留足够的旧类别样本进行联合微调。

如果你正打算做类似的项目,我建议不要一上来就追新模型,先把ResNet50这套经典组合跑通,再把数据、训练、推理、部署这条链路完整走一遍。模型可以换,但解决问题的思路和工程能力才是真正能沉淀下来的东西。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进