山里采的野花野草、药房抽屉里堆的干燥药材,很多新手是真的分不清。薄荷和荆芥叶形相似,金银花和连翘花黄艳艳一片,更别说那些切成片、晒成干之后长得几乎一模一样的根茎类药材。过去这类判断只能靠老药工的经验,现在用深度学习做图像分类,手机拍一张照就能给出一组候选类别和置信度。这篇文章聊的就是我做的一个基于PyTorch框架、EfficientNetV2模型的中草药识别分类系统,从数据集构造、模型训练、调参优化到最终部署,完整源码的落地全过程。想走一遍图像分类完整链路、或者想把自己的垂直场景做成识别工具的朋友,这篇应该能帮上忙。
1. 中草药识别任务定位与系统整体思路
1.1 这类系统的核心价值到底在哪
中草药识别看起来是个很小的分类问题,但放到真实场景里,需求非常明确而且迫切。药农采收、市场流通检查、野外采药科普、养生爱好者自拍识药,都需要一个低门槛的快速鉴别手段。传统做法靠的是药材图谱和专家经验,效率低且门槛高;而一个图像分类模型只要数据到位,就能在几秒钟内给出Top-5候选,配合置信度排序,完全可以把人工初筛的工作量降下来。
从技术角度看,中草药识别是一个标准的细粒度图像分类任务。细粒度意味着很多类别之间的差异并不是“猫和狗”那种根本性的差异,而是“陈皮和青皮”“白芷和当归”这种局部纹理、颜色、形状上的细微区别。这类任务对模型的特征提取能力要求更高,单纯堆层数未必有效,反而需要更高效的特征复用机制和更强的正则化手段,这也是我最终选定EfficientNetV2的重要原因。
1.2 识别对象与场景差异:鲜品、干品与饮片
做中草药识别,第一个要面对的问题是识别对象不一致。同一个类别,新鲜植物带花带叶、干燥整株、切成薄片的饮片,视觉特征差异非常大。刚开始我就踩过这个坑:只训练了新鲜叶片的数据,模型一到干燥饮片上准确率直接崩盘。
后来我把数据集按来源分成三类对待:
- 鲜品:植物叶片、花、果实的实拍图,环境光复杂,背景干扰多。
- 干品:干燥后的整株或部位,颜色整体偏黄褐,纹理比鲜品更明显。
- 饮片:炮制后的切片或碎块,形态高度相似,非常依赖纹理和切面特征。
最省事的做法是不区分来源,把三种形态的图片全部归到同一个类别下,让模型自己学习“形态不变映射”。数据量够大的时候,这种方式效果不错,而且代码简单。但如果某几类的鲜品图特别多、干品图特别少,模型就会对多数形态产生偏好。这时候可以考虑更细的做法:先用一个前置二分类模型判断“鲜品还是干制品”,再交给品种识别模型。我在项目里先用了混合数据集方案把主流程跑通,后续再按形态拆分做分支模型,两步走的方式更稳。
1.3 系统整体技术栈与流程
整个系统的技术选型围绕着“快速复现、清晰可改、最好能直接部署”这三个目标。整体流程是这样的:
图片采集 -> 清洗与去重 -> 同类别归并 -> 数据划分 -> 预处理 -> 模型训练 -> 评估分析 -> 推理预测 -> 接口封装技术栈方面:
- 框架:PyTorch,生态完整,调试方便,torchvision自带EfficientNetV2预训练权重。
- 模型:EfficientNetV2-S作为主力,兼顾参数量、精度和推理速度。
- 数据增强:torchvision.transforms加RandAugment、MixUp、CutMix等强增强手段。
- 工具:albumentations用于复杂增强,scikit-learn做数据划分和评估指标,matplotlib和TensorBoard做训练可视化。
这套组合的好处是每一环都有成熟方案,不需要自己造轮子。只要数据集准备好,几乎可以照着代码直接跑通。
2. 为什么选EfficientNetV2:模型选型的完整考量
2.1 EfficientNetV2相比V1改了什么
EfficientNetV2模型源自2021年发布的一篇论文,核心思路还是在效率上做文章。相比第一代EfficientNetV1,最大的变化有两点。
第一个变化是用Fused-MBConv替换了部分普通MBConv。MBConv是MobileNet风格的深度可分离卷积块,计算量小但特征复用能力较弱;Fused-MBConv把卷积核先融合成一个普通的3x3卷积,虽然单看FLOPs略高,但在早期stage上反而能显著提升训练速度和GPU利用率。EfficientNetV2的stage设计在浅层用Fused-MBConv,深层用带注意力机制的MBConv,兼顾了速度和精度。
第二个变化是渐进式训练策略(Progressive Training)。简单说就是训练初期图像分辨率小、正则化弱、数据增强轻,随着epoch推进逐步增大分辨率、增强正则化强度。这个策略让EfficientNetV2在ImageNet上的训练速度快了很多,收敛也更稳定。虽然我自己在训练时没有完全照搬渐进式策略,但把其中的思想用了一部分:前10个epoch用相对温和的增强,后期再开RandAugment和MixUp,效果确实更稳。
2.2 与ResNet、MobileNet、Vision Transformer的取舍
选模型之前,我做了几组对比。这里摘要整理成一张表,方便直观判断:
| 模型 | 参数量 | ImageNet Top-1大约值 | 单位精度成本 | 适用场景 |
|---|---|---|---|---|
| ResNet50 | 25.6M | 76%~79% | 中规中矩 | 经典基线,迁移学习稳定 |
| EfficientNetV1-B4 | 19M | 82%左右 | 训练较慢 | 精度好但训练耗时长 |
| EfficientNetV2-S | 21M | 83%~84% | 快,均衡 | 中草药识别主力选择 |
| EfficientNetV2-M | 54M | 85%左右 | 较慢 | 数据充足时的高精度选择 |
| MobileNetV3-Large | 5.4M | 75%左右 | 极轻量 | 移动端或嵌入式部署 |
ResNet50的优势是结构成熟、pretrained权重到处都有、调参经验丰富,但同等参数下它的Top-1精度在细粒度任务上明显弱于EfficientNet系。MobileNetV3适合端侧推理,但这种细粒度识别对它的特征表达能力来说偏吃力,而且中草药类别数一旦超过几十个,softmax输出空间的压力也会变大。Vision Transformer系列很强,ViT-S/16在ImageNet上可以到81%以上,但需要更大的数据量来喂,训练收敛也更慢,对于几千张中草药图片的项目并不划算。
2.3 对中草药数据集的实际适配性
中草药图片数据集普遍存在几个特点:类别多但每类数量少、类间相似度高、背景干扰大。EfficientNetV2-S对这类任务有几点天然优势:
- 参数量只有20M左右,不容易在小数据集上严重过拟合。
- SE注意力机制能自动关注药材主体区域,抑制背景干扰。
- 经过ImageNet预训练后,浅层和中层的通用特征迁移效果好,需要的微调数据量相对少。
实际跑下来的结果也印证了这一点。同一个中草药数据集上,ResNet50的Top-1准确率大约在86%左右,EfficientNetV2-S在相同训练设置下能到90%上下,参数量反而更少。在细粒度差异比较大的“陈皮/青皮”“金银花/连翘”这类近缘对上,EfficientNetV2的SE模块确实更容易学到判别性的局部特征。
3. 数据集构建与预处理:最容易被低估的环节
3.1 数据从哪里来、怎么清洗
模型能学到什么程度,数据说了算。中草药图片的公开数据集很零散,往往要靠自己采集和整理。我主要用了这几个来源:植物志图片库、常见中药图谱、公开的植物分类数据集子集,以及少量手工拍摄补充。
清洗是这步的重头戏。网络图片标签混乱,常见的几个问题:
- 图文不符:配图是整株植物,但标签写的是炮制饮片名。
- 重复图:同一张图在不同网站反复出现,不经过去重会严重过拟合。
- 混入无关内容:营销插图、漫画示意图、文字水印图片都混在数据里。
我的清洗流程是:先按类别文件夹浏览一遍,删除明显的错图;再用感知哈希相似度去重,两张图相似度高于0.9就删掉一张;最后把每个类别的图片数量统计出来,低于阈值50张的类别直接放弃或补充采集。这一步很费时间,但直接影响后续模型上限。
3.2 目录结构、划分策略与Dataset实现
数据集目录结构我建议直接按PyTorch的ImageFolder习惯组织:
herb_dataset/ ├── train/ │ ├── gouqi/ │ ├── jinyinhua/ │ ├── juhua/ │ └── ... ├── val/ └── test/用拼音或者英文目录名,避免中文路径在部分库和平台下出现编码问题。类别名称单独维护一个映射文件category.json,训练结束推理时再显示中文名。
数据划分要特别注意同一个来源的图片不能同时出现在训练集和验证集。我一开始偷懒直接用随机划分,结果验证集准确率虚高了两个点,原因就是网络来源的重复图同时落到了train和val里。后来改用按图片URL来源分组划分,或者先用感知哈希去重再划分,评价指标才恢复正常。这一条建议所有做图像分类的同学都记住。
Dataset实现可以直接用torchvision.datasets.ImageFolder,但为了后面加入更复杂的数据增强和控制采样逻辑,我改写了自定义版本。
from PIL import Image from torch.utils.data import Dataset class HerbDataset(Dataset): def __init__(self, image_paths, labels, transform=None): self.image_paths = image_paths self.labels = labels self.transform = transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img = Image.open(self.image_paths[idx]).convert("RGB") label = self.labels[idx] if self.transform: img = self.transform(img) return img, label3.3 数据增强方案设计与细节
中草药图片的多样性主要来自拍摄角度、光照、背景和形态差异。数据增强的目标就是模拟这些变化。我的增强方案分两档:基础增强和强增强。
基础增强:
| 增强方式 | 参数范围 | 作用 |
|---|---|---|
| RandomResizedCrop | scale 0.08~1.0 | 模拟距离远近和裁切 |
| RandomHorizontalFlip | 概率0.5 | 增加水平翻转不变性 |
| ColorJitter | brightness 0.3, contrast 0.3, saturation 0.3 | 适应不同光照环境 |
强增强在基础增强之上加入:
- RandAugment:随机选2种增强策略,幅度9级,包括旋转、平移、剪切、对比度、色彩均衡等。
- MixUp:按0.2的alpha从Beta分布采样融合系数,把两张图按比例混合,标签也按同样比例混合。
- CutMix:将一张图的部分区域剪切到另一张图上,标签按面积比例混合。
训练初期用基础增强,让模型先稳定收敛;中期开RandAugment和MixUp,提升泛化能力;最后一个阶段把增强强度降下来,让模型在接近真实分布的数据上微调。这种方法模仿了EfficientNetV2渐进式训练的思想,实测能大幅减缓过拟合。
4. 基于PyTorch的完整训练实现
4.1 模型定义:加载预训练权重并替换分类头
EfficientNetV2在torchvision里直接有预训练版本,选S规格即可。加载时把最外层的分类器替换成自己的输出维度,这是迁移学习最常规的操作。
import torch.nn as nn from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights class HerbClassifier(nn.Module): def __init__(self, num_classes, pretrained=True): super().__init__() if pretrained: self.backbone = efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1) else: self.backbone = efficientnet_v2_s(weights=None) in_features = self.backbone.classifier[1].in_features self.backbone.classifier = nn.Sequential( nn.Dropout(p=0.2, inplace=True), nn.Linear(in_features, num_classes) ) def forward(self, x): return self.backbone(x)注意最后原分类器是一个Linear层,直接替换时要用in_features获取原输入维度。Dropout保留0.2,这个值对细粒度分类任务比较稳妥,太大容易欠拟合。
4.2 训练循环与评估函数
训练循环里我重点加了混合精度AMP和梯度裁剪。AMP能用更少的显存跑更大的batch,梯度裁剪能防止训练后期loss突然爆炸。
import torch from tqdm import tqdm def train_one_epoch(model, loader, criterion, optimizer, scaler, device): model.train() total_loss = 0.0 correct = 0 total = 0 for images, labels in tqdm(loader, desc="Training"): images, labels = images.to(device), labels.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs = model(images) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() total_loss += loss.item() * images.size(0) _, preds = torch.max(outputs, 1) correct += (preds == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total评估函数要和训练函数彻底分离,带上model.eval()和torch.no_grad(),同时输出混淆矩阵和分类报告,方便后期分析。
import torch from sklearn.metrics import confusion_matrix, classification_report @torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() total_loss = 0.0 correct = 0 total = 0 all_labels = [] all_preds = [] for images, labels in loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) total_loss += loss.item() * images.size(0) _, preds = torch.max(outputs, 1) correct += (preds == labels).sum().item() total += labels.size(0) all_labels.extend(labels.cpu().numpy()) all_preds.extend(preds.cpu().numpy()) acc = correct / total cm = confusion_matrix(all_labels, all_preds) report = classification_report(all_labels, all_preds, zero_division=0) return total_loss / total, acc, cm, report4.3 迁移学习训练策略:先冻结后微调
刚上来就全参数微调有两个问题:一是模型在初期loss下降很猛但泛化差,二是解冻全部层之后训练不稳定。我用的是分阶段训练策略。
第一阶段冻结backbone全部参数,只训练新的分类头。此时规模很小,用稍高的学习率快速把分类头调到合理范围。
for param in model.backbone.parameters(): param.requires_grad = False optimizer = torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3 )第二阶段解冻backbone最后两到三个stage,用较低的学习率整体微调。浅层保留ImageNet学到的通用边缘纹理特征,深层面向中草药数据做适配。
for name, param in model.backbone.named_parameters(): if "features.5" in name or "features.6" in name or "classifier" in name: param.requires_grad = True else: param.requires_grad = False这个解冻位置需要根据backbone结构具体看。EfficientNetV2-S的features分6个stage,从features.5开始解冻是一个不错的起点。实际训练时我会打印每层的requires_grad确认一下,防止参数冻结配置错误。
5. 训练配置与调参经验
5.1 超参数选择与学习率策略
训练是否成功,一半看超参。我的推荐配置如下:
| 超参数 | 设置 | 说明 |
|---|---|---|
| 训练图像尺寸 | 224x224 | 兼顾速度和精度,EfficientNetV2原生支持 |
| 验证图像尺寸 | 224x224 | 与训练保持一致 |
| Batch Size | 32 | 显存不够时可以降到16 |
| 优化器 | AdamW | 对权重衰减的处理比Adam更稳 |
| 分类头学习率 | 3e-4 | 第一阶段用 |
| backbone学习率 | 1e-4 | 第二阶段微调用 |
| 学习率调度 | CosineAnnealingLR | T_max设为剩余epoch数 |
| Epoch | 40~60 | 数据量少时建议40,数据充足可到60 |
| 标签平滑 | 0.1 | 防止输出过于自信 |
| 权重衰减 | 1e-4 | 抑制过拟合 |
学习率设置上,AdamW加上CosineAnnealing的组合比固定学习率效果好很多。学习率在训练后半段cosine式下降,能够让模型在loss plateau区域更细腻地收敛,不容易来回震荡。第一阶段结束后,把backbone解锁时学习率直接降到1e-4,而不是继续用3e-4,这是很多新手最容易忽略的。
我监控过几次loss曲线,一个典型特征是:若第一阶段结束时验证集准确率还在快速上升,说明训练epoch不够;若训练集准确率已经接近100%而验证集停滞,说明增强强度不够或模型容量过剩,需要加强正则化。
5.2 类别不均衡与损失函数处理
中药类别天然不均衡。像枸杞、菊花这种大众药材,网上图片量很大;某些冷门药材可能只能找到三五十张。直接拿原始分布训练,模型会偏向样本多的类别,少数类别精确率惨不忍睹。
最简单的处理方式是把类别权重直接放进CrossEntropyLoss。
class_weights = torch.tensor(weights, dtype=torch.float32).to(device) criterion = nn.CrossEntropyLoss(weight=class_weights)权重计算方法有很多种,我用的是有效样本数公式:weight = (1 - beta) / (1 - beta^n),其中beta取0.999,n是当前类别样本数。这个公式比简单的“1/样本数”更平滑,不会让少数类权重过大导致训练震荡。
更进阶的处理是用Focal Loss。它在CrossEntropy基础上增加了调制因子,让模型更关注难分类样本,对中草药近缘品种非常有效。我在训练后期试过把损失从CrossEntropy换成Focal Loss,少数类的F1分数提升了约3个百分点。如果只是想快速跑通,先用加权CrossEntropy就够了。
5.3 训练过程的监控与可视化
训练状态不能只看终端里打印的loss数字,我并行用了两个工具:TensorBoard和matplotlib绘制的准确率曲线图。
TensorBoard主要看四个指标:训练loss、验证loss、训练准确率、验证准确率。两两对比就能快速判断过拟合程度和训练状态。例如训练loss持续下降但验证loss在第20个epoch开始上升,说明进入过拟合区,需要提前终止或增强正则化;验证loss下降但验证准确率不上升,说明数据划分或标签存在问题。
matplotlib曲线图用于复盘。我把每次实验的模型配置、超参数、增强策略记录到一组CSV里,再统一绘制对比曲线。这个习惯帮我省了大量时间:想看看某个增强手段是否有效,直接画图比对即可,不需要靠记忆。
6. 评估指标体系与实测结果分析
6.1 除了Top-1准确率,还要看哪些指标
很多入门玩家只盯着Top-1准确率,这在普通分类问题里够用,但在中草药识别这种细粒度场景里,Top-1会掩盖大量问题。我实际看的是Precision、Recall和F1-Score这些更细的指标。
举个实际遇到的例子:整个模型的Top-1准确率有90%,但里面某个冷门药材类别的Recall只有40%,也就是说模型把六成这类药材都判成了其他类。如果只看平均准确率,这个问题完全不会被发现;一旦到真实场景使用,这类的漏检率是高得没法接受的。
使用scikit-learn的classification_report可以一次性拿到所有类别的精确率、召回率和F1分数。我一般固定检查两类问题:
- 样本数最多的几个类别:如果它们的Precision高但Recall低,说明模型“宁可漏判也不错判”,需要调整决策阈值。
- 样本数最少的几个类别:如果Recall大幅偏低,优先做数据增强和类别权重调整。
6.2 混淆矩阵的实际使用场景
混淆矩阵比准确率图更有诊断价值。我打印出混淆矩阵热力图之后,可以一眼锁定经常被搞混的类别对。最常见的混淆发生在形态极其相似的品种之间,比如金银花和连翘、白芷和当归、薄荷和荆芥。
针对混淆对的处理办法:
- 收集更多该类别的数据,尤其是容易混淆的形态。
- 单独评估这对类别的区分度,如果Acc接近50%,说明特征信息不够。
- 检查标注是否错误。我遇到过一种情况:某两类始终混淆,结果排查数据发现是采集阶段标签就贴反了,浪费了两天训练时间。
混淆矩阵还能帮我看清“模型到底学了什么”。比如模型频繁把“干枸杞”判成“干山楂”,说明它关注的是颜色和表面的皱缩纹理,还没学到形状上的差异,可以通过增强训练数据或调整裁剪区域来改善。
6.3 单张图片推理与可视化测试
训练完成后要做的第一件事,就是拿单张图片做推理,把整个pipeline跑通。这个步骤简单但容易被忽略,很多问题都是在“模型已经训练完了,但部署时才发现预处理不一致”后才暴露出来的。
推理函数里最关键的坑是预处理务必和训练保持一致。训练时用RandomResizedCrop,推理时必须用Resize加CenterCrop;训练时做了Normalize,推理时也必须做完全相同的Normalize,均值标准差用ImageNet预训练的标准值即可。
from PIL import Image from torchvision import transforms def predict_image(image_path, model, class_names, device="cuda"): 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]), ]) img = Image.open(image_path).convert("RGB") x = transform(img).unsqueeze(0).to(device) with torch.no_grad(): logits = model(x) probs = torch.softmax(logits, dim=1) topk_idx = torch.argsort(probs, descending=True)[0] return [(class_names[i], probs[0][i].item()) for i in topk_idx[:5]]输入图片时还要检查尺寸。真实拍摄的药品图很多带有大块背景,直接把全图缩到224会丢失细节。我实测先把图片按最长边缩放到256,再中心裁剪224,比直接resize效果更好。如果想更精细,可以先在前置程序里做一个目标检测或简单的主体裁剪,把药材区域抠出来再送进分类器。
7. 典型问题排查与实战避坑
7.1 模型过拟合与数据增强踩坑
中草药数据集规模普遍不大,过拟合几乎是必然会遇到的问题。我的症状很典型:训练集准确率很快就到了97%,验证集卡在85%附近上不去。
排查思路和解决办法,按优先级排序:
- 降低模型容量:先把EfficientNetV2-S换成一个更小的配置试试,如果验证集准确率反而上升,说明当前模型对数据量来说容量过大了。
- 增强正则化:提高Weight Decay到1e-3,Dropout从0.2调到0.3,训练后期加入更大的RandAugment幅度。
- 添加MixUp和CutMix:这两个增强对细粒度分类帮助极大,但要注意MixUp会让训练loss变高,这时候不要慌,验证集准确率才是真正需要关注的指标。
- 提前终止:验证集准确率连续5个epoch不涨就停止训练,节省时间也防止过度拟合。
还有一种隐蔽的过拟合:数据划分时没做去重,训练集和验证集里有相同来源的重复图。这种情况下验证集准确率虚高,模型部署到真实场景立刻打回原形。建议做严格去重后再划分,甚至可以考虑用图片特征做聚类式划分。
7.2 显存不足、训练速度慢的解决办法
EfficientNetV2-S输入224分辨率时显存压力不算大,但如果你用更大分辨率或超大batch,还是容易遇到OOM。我碰到一次显存16GB的设备跑batch size 64直接报错,换到batch 32才稳定。
显存不足的常规解法有几种:
- 降低batch size,配合梯度累积补偿。
- 开AMP混合精度,显存直接砍一半。
- 减少输入图像尺寸,从224降到192甚至160,对中草药这种小目标分类影响有限。
- 检查有没有无意中把多个中间变量保留在显存里,比如
outputs = model(images)之后,如果后面还用outputs和labels计算loss,中间变量就会保留;用完立即释放。
训练速度慢的问题,我主要靠两招解决:一是全流程用AMP,在图像分类任务里通常能带来1.5到2倍加速;二是DataLoader的num_workers设置到CPU核心数的两倍左右,pin_memory=True,在Linux系统下数据读取瓶颈基本消失。
7.3 识别混淆严重怎么办
如果混淆矩阵里固定有几对类别分不开,先不要急着重训模型。我的排查顺序是:
第一检查数据标注。把两个混淆类别的样本分别随机抽50张,人工过一遍,确认是否真的标注正确。我在项目里就靠这个操作挽救了一整轮训练,发现将近20张图片是从网上爬来的错误标签。
第二对比特征可视化。把两个类别的图片输入模型,提取分类头前一层输出的特征向量,用PCA或t-SNE降维可视化。如果两个类别的点在特征空间完全重叠,说明模型没有学到区分性特征;如果点群分得开但仍有大量错误样本,则可能是决策边界问题,可以尝试加Margin Loss。
第三针对性补数据。对混淆对专门采集更多样本,特别是容易搞混的干品形态。这是最直接也最有效的手段。
8. 部署与扩展方向
8.1 导出TorchScript和ONNX模型
训练完成后,下一步就是把PyTorch模型变成一个可以脱离训练环境运行的推理单元。我通常导出两种格式,按部署场景选用。
TorchScript适合PyTorch依赖链内的服务端部署,导出方式:
model.eval() example = torch.randn(1, 3, 224, 224).to(device) traced_script_model = torch.jit.trace(model, example) traced_script_model.save("herb_classifier.pt")ONNX格式则更通用,后端推理可以脱离Python环境。导出需要安装torch.onnx依赖,并指定opset版本。
torch.onnx.export( model, example, "herb_classifier.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}, opset_version=12 )导出后建议用ONNXRuntime做一次推理测试,对比PyTorch原始模型输出,通常会有微小浮点差异,但Top-1结论不应该改变。如果差异明显,大概率是导出时没把模型切到eval模式,或者BatchNorm层的统计量没有被固定。
8.2 Flask接口封装与调用示例
部署最直接的方式是用Flask封装一个HTTP接口,手机端或前端小程序直接上传图片即可识别。这是一个极简但能用的接口实现。
from flask import Flask, request, jsonify app = Flask(__name__) model = load_model("herb_classifier.pt") class_names = load_class_names("category.json") @app.route("/predict", methods=["POST"]) def predict(): file = request.files.get("image") if not file: return jsonify({"error": "no image"}), 400 img = preprocess(file.read()) with torch.no_grad(): logits = model(img) probs = torch.softmax(logits, dim=1) topk = torch.argsort(probs, descending=True)[0][:5] result = [{"name": class_names[i], "score": round(probs[0][i].item(), 4)} for i in topk] return jsonify({"results": result}) if __name__ == "__main__": app.run(host="0.0.0.0", port=5000)接口上线前一定要做图片大小和格式校验,防止超大分辨率图片拖垮服务。我这边对最大上传尺寸做了限制,常见图片类型只允许JPEG和PNG,并且对上传文件做Base64解码前先判断文件头,避免无效请求。
8.3 可扩展方向
这个系统做完之后,可以往几个方向继续扩展。
大模型方向的改进是用EfficientNetV2-S做特征提取器,接一个Metric Learning的embedding头,把分类任务转成度量学习任务,这样模型对未见过的中药材类别也有一定的开放识别能力,而不是只能在训练集类别里打转。
另一个实用扩展是主干模型加目标检测。分类模型默认图片里只有一个主体,但真实场景一张照片里可能同时出现多种药材。先在图片上检测出每个目标区域,再逐区域送入分类器,整个系统的实用性会提升一大截。
移动端部署是更贴近使用场景的方向。把ONNX模型转到TFLite或NCNN格式,就能在手机App里跑本地推理,不需要上传服务器,对隐私和网络要求都友好不少。
我在实际项目中还遇到过一个需求:用户希望能识别“哪味中药”之外,还希望看到该药材的性味归经和注意事项。这类知识库型需求,技术上就是分类结果去匹配后端一个药材信息表,属于很轻量的增量功能,但对产品完整度提升很明显。
最后想说的是,图像分类模型本身不难,难的是数据和场景的适配。做中草药识别尤其要关注形态差异、类别相似性和数据质量这三关。我自己在这套项目上最大的体会就是:花在数据清洗和评估分析上的时间,远比花在调模型结构上的时间来得值。遇到识别效果不理想,先别急着换网络结构,把混淆矩阵打开,看看模型到底在哪一类上翻车,往往比盲目调参有效得多。