简介基于水位的机器学习水瓶图像分类项目资料包面向机器学习初学者、图像分类实践者及计算机视觉课程设计人群用于解决如何根据水瓶中水位状态满水位、半水位或溢出自动对图像进行分类的问题。压缩包共488个文件大小约64.94MB其中429个jpeg与57个png构成图像数据集涵盖多种水瓶大小、形状与拍摄角度且状态清晰适合CNN训练与验证另有1个py脚本与1个ipynb Notebook提供数据加载、模型构建、训练评估等完整流程Notebook中采用ResNet50架构完成图像分类实验。已有185人学习下载。获取后即可在运行环境中复现水位分类结果也可基于自带数据与源码进行二次训练、调参与扩展类别对于理解卷积神经网络的实际应用、完成课程作业或毕业设计均有较高参考价值。1. 基于水位的机器学习水瓶图像分类是什么一个CNN任务不只是“把瓶子分个类”“基于水位的机器学习水瓶图像分类”这个标题核心不是区分矿泉水瓶或可乐瓶而是根据水瓶里的水位——空、低、中、高、满——对照片分档。它把CNN卷积神经网络和数据集、Jupyter Notebook、训练源码打包成一条完整pipeline输入一张水瓶照片输出水位属于哪一档。这类能力在灌装质检、智能饮水、实验容器管理里都有真实位置。适合刚入门深度学习的读者把图像分类从概念变成自己电脑上能跑通、能改参数、能看曲线的东西。先说明一下做水位分类比做普通瓶身分类更依赖数据一致性这恰恰是后面所有环节的地基。2. 从照片到数据集水位分类的数据组织与预处理规范2.1 先定“分类语义”再动手拍图标注不要一上手就拍一堆照片然后丢给CNN。水位分类的第一步是确定你到底分几类、边界在哪里。常见做法是分成五类empty水位线在瓶高5%以下、low5%-30%、medium30%-60%、high60%-90%、full90%-100%。也可以只分三类低于三分之一、三分之一到三分之二、高于三分之二。类数越少标注越一致模型越容易收敛。分类档位的选择直接决定后面的数据集规模五分类比三分类需要明显更多的样本因为相邻档位比如medium和high在照片上的差异本来就很小分得太细等于自己给自己增加难度。另一个常见疑问是水位明明是连续值为什么做成分类而不是回归回归思路看着更自然直接输出水位百分比但回归对拍摄角度和瓶子底部偏移非常敏感相机偏几度像素误差就会被放大成水位误差。分类只要保证档位边界不混乱即可对微小偏移有天然容忍度。这也是标题写“分类”而不是“检测”的原因。实际项目里如果需求只是“要不要加水”三分类比五分类更实用能显著减少标注和模型训练的压力。分类语义定了拍摄约束才能提出来。图像分类里CNN会同时学习“水位”和“周围一切”。如果训练照片里的瓶子永远是干净的透明塑料瓶、白色背景、左侧光源模型很可能把“右侧有阴影”当成“高水位”的线索。所以拍摄规范至少包含四点同一型号的瓶身、固定拍摄距离和角度、统一背景和光源、每个水位至少拍摄多张不同角度与光线偏移的照片。宁可照片总数少一点也要保证类别内差异足够。数据规模随之明确这类项目通常在自己桌面或工位拍摄几百张到一两千张是合理规模不要指望拿几十张图就能训出一个稳定模型。标注我建议用“按文件夹归类”而不是维护CSV因为torchvision的ImageFolder会按子目录名自动读取标签省掉手工维护表格的麻烦。提示拍摄前先把类别定义打印出来贴在拍摄位置标注的时候每拍一张都对照定义判断能明显减少后期发现标注不一致返工的情况。2.2 目录结构与Dataset类用ImageFolder把文件夹变成能训练的Loader最小可运行的数据集组织长这样train和val下面分别是五个类别子目录每个子目录放对应水位档位的照片data/ ├── train/ │ ├── empty/ │ ├── low/ │ ├── medium/ │ ├── high/ │ └── full/ ├── val/ │ ├── empty/ │ ├── low/ │ ├── medium/ │ ├── high/ │ └── full/ └── test/ └── ...如果只想调参train和val两个目录就够test目录留到最后验货用。目录建好之后加载代码非常短import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 验证集只做统一尺寸和归一化绝不做数据增强 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]) ]) train_data datasets.ImageFolder(rootdata/train, transformval_transform) val_data datasets.ImageFolder(rootdata/val, transformval_transform) train_loader DataLoader(train_data, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_data, batch_size32, shuffleFalse, num_workers2) print(train_data.classes) # 五个类别名按字母序排列 print(train_data.class_to_idx) # 类别名到标签编号的映射 print(len(train_data), len(val_data))逻辑说明ImageFolder会自动把“data/train”下每个子目录当作一个类别子目录名就是标签不需要额外维护CSV文件。.classes给出按字母序排列的类别名.class_to_idx给出类别名到整数标签的映射这两行打印能帮你确认Dataset没有读错目录结构。参数说明train_loader里shuffleTrue保证每个epoch训练样本顺序被重新打乱避免模型学到批次顺序val_loader里shuffleFalse让评估输出顺序和真实标签一一对应方便后面画混淆矩阵。batch_size32在图片224×224、RGB三通道时单卡显存占用不算大如果机器显存只有4G直接降到16或8。num_workers在Windows的Notebook里建议设成2开太大重新创建worker时容易报错这个坑第5章会专门讲。2.3 预处理与数据增强先固定拍摄条件再谈增强数据增强有个铁律训练集上用验证集和测试集上绝对不用。上面代码里val_transform是一份“纯净”的变换只做Resize、ToTensor、Normalize保证验证结果稳定可比较。训练集在此基础上再叠加增强操作train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.85, 1.0)), transforms.ColorJitter(brightness0.15, contrast0.15), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])为什么有意不加RandomHorizontalFlip水位分类和猫狗分类在这里有本质区别水平翻转会让瓶身上的标签、刻度、文字镜像分类器很容易去学“文字方向”这种和水量无关的假特征。如果你拍摄的瓶子完全没有文字和刻度可以加翻转否则别加。RandomResizedCrop的scale取0.85~1.0而不是PyTorch默认的0.08~1.0是为了避免裁掉大部分瓶身。水位判断依赖水位线和瓶底的相对位置裁太多等于把答案撕掉。归一化的mean/std直接沿用ImageNet统计值后续如果做迁移学习这套参数保持一致即可如果自建CNN从零训练可以按实际数据集的均值和标准差重新计算。最后的建议是增强只是兜底不是救火。固定光源、固定背景、固定拍摄设备比任何增强策略都重要。用手机拍没问题但全程用同一台手机、同一个拍摄距离不要中途换设备或者突然开闪光灯。这一条价值很大因为很多水位分类项目最终效果不理想回头查问题都出在数据一致性上。3. CNN选型与模型搭建分清“从零训练”和“迁移学习”的边界3.1 数据量决定模型方案为什么水位任务默认走迁移学习CNN处理图像的逻辑是图像本质上是网格状的像素矩阵卷积核在局部窗口滑动提取边缘、纹理、颜色过渡这些局部特征然后通过池化逐步聚合出全局信息。水位分类恰好落在CNN的舒适区里——水位线是一个明确的边缘特征瓶身和水的颜色对比又是强线索CNN不需要复杂的注意力机制就能抓住这些模式。但CNN有一个现实前提它要学出稳定特征需要有足够多的样本。于是这里出现两条完全不同的路线自己搭一个几层卷积的小网络从随机权重开始训练或者直接加载ImageNet预训练模型、只替换输出层做迁移学习。表格里的选择逻辑基本可以套用数据规模推荐方案原因每类不足100张迁移学习冻结backbone只训练最后一层训练参数量小不容易过拟合每类100~300张迁移学习微调backbone后半段预训练特征能被适配到水位任务每类300张以上自建CNN或迁移学习都可行自建网络更可控推理更快对“基于水位的机器学习水瓶图像分类”这种典型小数据集项目自建CNN从头学习边缘、颜色、纹理需要大量样本训练样本少时准确率上不去预训练模型已经具备通用视觉特征只需要让最后的全连接层学会把“特征”映射到“水位类别”。实际落地时迁移学习几乎是默认选择。3.2 最小可训练的自建CNNConv2d BatchNorm MaxPool的经典组合如果你更想理解CNN内部机制或者数据量确实足够可以从这个三层小网络开始import torch.nn as nn class WaterLevelCNN(nn.Module): def __init__(self, num_classes5): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Dropout(0.3), nn.Linear(64, num_classes), ) def forward(self, x): return self.classifier(self.features(x))逻辑说明输入224×224×3的RGB图像每经过一次MaxPool2d特征图边长减半224→112→56→28。最后一层卷积输出64个通道AdaptiveAvgPool2d((1, 1))把任意尺寸的特征图压成1×1再Flatten成64维向量最后由Linear输出5个类别的得分。BatchNorm放在卷积和ReLU之间能显著加速收敛Dropout放在全连接层之前用于缓解过拟合。参数说明kernel_size3是性价比最高的选择比5×5参数量小比1×1能获取更多局部信息padding1保证卷积不改变特征图尺寸让池化层独立控制尺寸变化。如果手头只有两三百张数据建议把第一层输出通道从16降到8或者把Dropout比例提到0.5否则训练误差很快趋近于0、验证误差不降过拟合会来得非常快。3.3 迁移学习用预训练模型替换输出层保留视觉主干实际训练最快见效的代码就是这么几行import torch.nn as nn import torchvision.models as models num_classes 5 model models.resnet18(pretrainedTrue) # 冻结backbone只让最后一层全连接层可训练 for param in model.parameters(): param.requires_grad False model.fc nn.Linear(model.fc.in_features, num_classes)逻辑说明pretrainedTrue加载的是在ImageNet上训练过的权重此时模型已经能识别通用的边缘、纹理、物体部件。for param in model.parameters()把所有参数设为不可训练训练时反向传播不会更新这些层只有最后替换的model.fc会学习。这样训练参数量从几百万骤降到几十万数据量小也能稳住。参数说明冻结版适合每类不足100张的数据。model.fc.in_features是ResNet18全连接层的输入维度用这个字段做替换不需要硬编码数值。如果数据量超过每类200张可以解冻backbone的后半段做微调fine-tune常见做法是把学习率调小到原来的十分之一比如从1e-4改成1e-5避免破坏已经学好的底层特征。预训练模型的选择按场景来ResNet18轻量、推理快对水位分类足够MobileNetV3更省资源但训练收敛略慢ResNet50精度上限更高但小数据集上更容易过拟合不建议一上来就用。输出层接Linear(num_classes)就好不要手动加Softmax——CrossEntropyLoss内部已经包含Softmax计算训练时喂原始logits即可。4. Notebook里的训练闭环从数据加载到loss曲线的完整pipeline4.1 构建Notebook之前先固定随机种子和计算设备标题里的Notebook通常就是Jupyter Notebook用它跑深度学习项目很合适因为每一步的中间结果、loss曲线、错误样本图片都能即时看到。但在写训练循环之前先把“随机性”和“设备”处理成固定值。深度学习的随机性来源包括数据加载顺序、权重初始化、Dropout、数据增强不固定seed会导致同样代码两次跑出的结果不一样后面调参时你根本分不清提升是来自代码修改还是运气。import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) set_seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu) print(device)逻辑说明torch.manual_seed控制PyTorch的CPU随机数生成器torch.cuda.manual_seed_all控制所有GPU设备。random和numpy的seed也要一起固定因为数据增强里的随机操作可能用到它们。device统一写在Notebook最顶部后面所有cell都引用这个变量换机器时只需要改这一处。参数说明seed取42是习惯取其他任意整数效果一样关键是固定。如果用的是Windows且DataLoader开了多进程还需要在DataLoader里加上generatortorch.Generator().manual_seed(seed)否则worker的随机性无法完全控制。4.2 训练循环与超参数把关键参数一次设定清楚Notebook适合把训练逻辑拆成函数再写循环。完整的训练闭环一般包含三个部分损失函数、优化器、训练和评估函数。以下是可以直接粘贴到Notebook cell里的版本import torch.nn as nn criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max25) def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss, correct, num 0.0, 0, 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) correct (outputs.argmax(1) labels).sum().item() num images.size(0) return total_loss / num, correct / num def evaluate(model, loader, criterion, device): model.eval() total_loss, correct, num 0.0, 0, 0 with torch.no_grad(): 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) correct (outputs.argmax(1) labels).sum().item() num images.size(0) return total_loss / num, correct / num epochs 25 for epoch in range(1, epochs 1): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fepoch {epoch:02d} | train loss {train_loss:.4f} acc {train_acc:.4f} | val loss {val_loss:.4f} acc {val_acc:.4f})逻辑说明criterion用CrossEntropyLossPyTorch会在内部完成Softmax和损失计算所以模型输出直接传原始logits就行。optimizer.zero_grad()把上一步的梯度清零loss.backward()算出当前batch的梯度optimizer.step()更新参数这是PyTorch训练循环的标准四步。参数说明AdamW相比原始Adam权重衰减更规范对这类任务默认优先选它。lr1e-4适合迁移学习或微调场景如果是自建CNN可放宽到1e-3。weight_decay1e-4是小数据集上稳妥的L2正则强度。CosineAnnealingLR的T_max25和epoch数保持一致让学习率沿余弦曲线平滑降到接近0通常比固定学习率收敛更稳定、末尾效果更好。model.train()和model.eval()必须成对出现因为它们分别控制BatchNorm和Dropout的行为BatchNorm在训练模式下用当前batch的均值方差在评估模式下用累积统计量切错会让验证结果忽高忽低。with torch.no_grad()在评估时关闭梯度计算省显存也更快。4.3 评估与可视化准确率之外要用混淆矩阵定位错误准确率会掩盖错误模式。比如medium和high相邻档的混淆和empty与full的混淆背后的原因完全不同。前者是视觉上本身就接近后者是模型没学好特征。因此训练结束后一定要画混淆矩阵import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix, classification_report import numpy as np all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) preds model(images).argmax(1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds, target_namestrain_data.classes)) plt.figure(figsize(6, 5)) plt.imshow(cm, cmapBlues) plt.xticks(range(len(train_data.classes)), train_data.classes, rotation45) plt.yticks(range(len(train_data.classes)), train_data.classes) plt.colorbar() plt.xlabel(Predicted) plt.ylabel(True) plt.show()逻辑说明argmax(1)取模型输出得分最高的下标作为预测类别cpu().numpy()把tensor转成numpy数组方便交给sklearn统计。confusion_matrix返回一个5×5矩阵第i行第j列表示“真实类别i被预测为类别j”的样本数。classification_report里的precision精确率和recall召回率值得单独关注。对水位任务最该看的是相邻类别之间的混淆如果medium样本大量被判成high说明数据采集时类别边界控制得不好或者拍摄目标没有对准类别中心。这一步完成后再回到5.4节的解决方案里调整数据而不是盲目加训练轮数。5. 避坑水位分类最不值得踩的5个翻车现场5.1 同一个瓶子拍到黑换瓶子后模型直接失效现象训练集和验证集共用一只透明水瓶验证准确率很高一旦在实际环境换成另一只瓶型或壁厚不同的水瓶准确率瞬间掉20个点甚至更多。原因CNN会把瓶身的透明纹理、塑料折射、口径形状这些和“这个瓶子”绑定的外观特征一并学进去。水位分类真正该学的是“水位线和容器底部的相对位置”但网络没有这个先验它只会从像素里找最容易区分的线索瓶型特征往往比水位线更显眼。解决拍摄时至少使用3只不同款式的瓶子让每个水位档位都在每只瓶子上出现。验证集里放一只训练时从未见过的瓶子如果这只瓶子的准确率明显偏低说明模型在识别瓶子而不是水位。这类问题只能在数据采集阶段解决靠增强和调参都治不了根。5.2 类别不均衡空瓶和满瓶容易分中间档样本稀缺现象empty和full类的准确率很高low和medium的召回率低得可怜整体准确率看着还行但中间档每次都被边缘化。原因拍摄者拍“空”和“满”容易判断也容易执行中间水位拍起来费劲导致样本量天然不平衡。CrossEntropyLoss在类别不均衡时优化目标会偏向样本多的类别少数类怎么学都学不好。解决首先检查每类样本数量。可以用train_loader.dataset.targets统计标签分布如果样本数量差距超过2倍给损失函数加类别权重weights torch.tensor([...])然后nn.CrossEntropyLoss(weightweights.to(device))权重按各类样本数的倒数设置。另一个更直接的做法是WeightedRandomSampler在采样阶段把少数类样本重复采上来两者选一种即可。5.3 过拟合训练集98%验证集只有78%现象训练准确率一路走高接近100验证准确率在第10个epoch后不再上升甚至开始下降train loss和val loss之间的差距持续拉大。原因模型参数远多于样本信息量最后把训练集的背景、影子、反光这些细节都背下来了。水位分类项目数据量小撞上过拟合几乎必然只是时间早晚。解决先换更小的模型比如把ResNet18换成更窄的自建CNN或者给全连接层把Dropout从0.3提到0.5。数据增强可以加强一点比如加RandomRotation(5)但注意不要加水平翻转。最后用早停兜底验证集准确率连续5个epoch不再提升就保存当前模型并结束训练不要傻傻跑完预设的所有epoch。5.4 相邻类别被判混medium和high在照片上就是很难区分现象混淆矩阵里medium样本集中落在high列或者low和medium相互混杂错误不是完全发散而是系统性向相邻类别偏移。原因这类任务很难有绝对边界。同一张瓶子照片水位在60%还是65%人眼都有可能判错何况网络。如果标注时把medium定义为30%-60%、high定义为60%-90%两张接近60%的照片一个标medium一个标high模型学到的边界就是模糊的。解决拍摄时不要卡着边界取样。五分类任务建议以25%、50%、75%为目标水位去拍让类别中心明显拉开也可以直接把档位合并成三类减少天然难分的邻接档。另外可以在推理阶段引入“置信度低于阈值就待人工确认”的策略这个第6章展开讲。5.5 Notebook环境坑kernel崩溃和数据重复加载现象训练到一半kernel直接死掉或者每个epoch开始前都要等十几秒训练完成后又在某个cell里重新加载了一次数据集。原因这类坑多半是DataLoader的num_workers开得过大Windows下spawn出的多个worker同时拷贝图片内存被撑爆后kernel被杀另一个是Notebook里的cell执行顺序混乱数据集加载cell被反复执行。解决Windows下num_workers设成0或2不要贪多。把Dataset和DataLoader的构建放到Notebook最顶端的cell里只执行一次后续训练循环引用同一个loader对象。kernel死了之后优先重启kernel、按顺序从头执行所有cell不要跳着执行。我习惯把Notebook分成“数据准备、模型结构、训练函数、训练循环、可视化”五个区块每个区块开头的cell独立可重跑避免漏掉某个初始化状态。6. 进阶从“能跑”到“能用”——置信度阈值与错分样本复盘模型在验证集上做到90%准确率离“能投入使用”还差两步。第一步是把argmax换成带置信度的判断。argmax只返回得分最高的类别模型对两个档位都拿不准时也会硬着头皮给一个答案。而Softmax输出的概率分布能告诉你它到底有多确定probs torch.softmax(model(image), dim1) confidence, pred probs.max(1) final_label pred if confidence.item() 0.6 else -1 # -1 表示待人工确认逻辑说明softmax把原始logits转成总和为1的概率分布max(1)同时取出最高概率和对应类别。当最高置信度低于0.6时不输出具体档位而是交给人工确认。阈值设置看应用场景流水线上宁愿多拦截也不能漏放阈值调到0.7甚至0.8只是给用户显示一个水位提示0.5就够。注意阈值需要用验证集去标定而不是拍脑袋定。第二步是错分样本复盘。把验证集里预测错误且置信度较高的图像按真实类别汇总每类挑出5张拼成一张图肉眼观察。如果错的是medium和high多半是标注边界问题如果错的是空瓶和满瓶多半是拍摄光照问题。我自己的习惯是训练完后一定会看一遍错图虽然这个环节不直接提升指标但它能暴露出数据集的结构性问题后续决定要不要扩数据、加类别、换模型时这套错图比训练曲线更有说服力。训练源码、评估脚本和Notebook最终都会沉淀为这份pipeline的一部分下次换一个容器场景只需要换数据、改类别数、重跑一遍。这个方向投入产出比不低只要数据一致性做到位剩下的都按流程走。希望帮到你。本文还有配套的精品资源点击获取