简介面向图像分类与刀具磨损检测任务的微铣削刀具磨损图像数据集适合正在学习YOLOv5分类、CNN图像识别或需要可直接训练的数据集的学生与开发者使用。资源已按训练集、测试集划分共两类样本对应好、坏等磨损状态具体类别映射见JSON类别字典文件训练集3600张、测试集900张可配合YOLOv5分类工程或自建分类网络直接训练目录下按类别文件夹存放接入数据加载器时几乎不需要额外整理。包内共2000个文件以1998张JPG图片为主另含1个Python可视化脚本和1个JSON字典压缩包大小165.2MBshow脚本可随机展示图片方便快速核对标注与划分是否合理。已有502人学习下载适合需要快速获取规范划分数据、完成磨损状态识别实验或复现分类流程的入门及进阶用户。1. 微铣削刀具磨损图像识别一份划分好、能直接开训的图像分类数据集做微铣削加工的人都有这个痛点刀具磨损直接决定工件表面质量换刀不及时一批零件就废。主流做法是拍刀面图像交给图像分类模型判断磨损等级省去停机目测。但这个场景数据量小、类别边界模糊更麻烦的是很多数据集只给原始图片train/val/test 划分要自己动手类别字典还得对着文件夹名手搓。这份资源把这两件烦心事做完了磨损图像按类别分好、按文件夹保存训练集、验证集、测试集划分现成还带一份类别字典 JSON 文件把类别索引和磨损等级一一对应。用 torchvision 的 ImageFolder 可以直接加载省掉整理数据的功夫把精力放在模型选型和调参上。适合两类人做刀具状态监测、想快速跑出磨损识别基线的工程师拿小样本图像分类练手、需要干净标准数据集的算法同学。下面从数据集结构、JSON 字段、训练参数到避坑点完整拆一遍。2. 数据集与类别字典吃透文件夹结构和 JSON 字段再开训拿到资源先别急着训练。文件夹保存加类别字典 JSON 这套组合看着简单但里面有两处映射关系必须提前盘清楚文件夹名到整数索引的映射、整数索引到磨损等级名称的映射。搞反了或者没对齐训练再久结果都是错的。2.1 文件夹保存的目录结构为什么说它天生适配 ImageFolder这份数据集按「文件夹保存」组织目录层级是标准的分类数据集约定micro_milling_wear/ ├── train/ │ ├── normal_wear/ │ │ ├── 001.jpg │ │ └── ... │ ├── slight_wear/ │ ├── moderate_wear/ │ └── severe_wear/ ├── val/ │ └── ... └── test/ └── ...每个磨损等级一个子文件夹文件夹名就是类别名。torchvision 的datasets.ImageFolder天然按这个结构加载扫描根目录下所有子文件夹每个子文件夹自动成为一个类别生成(image_path, class_index)样本对。你唯一需要记住的是ImageFolder 按文件夹名的字母序来分配索引不是按文件夹里的排列顺序也不是按你脑子里「从轻到重」的语义顺序。这个设计对微铣削磨损场景很友好。磨损等级通常分三到五类每类一个文件夹直观、好维护。相比 CSV 标注或 JSON 标注文件夹方式最大的好处是增量加图不用改标注文件——现场新拍一批图像丢进对应文件夹就能参与下一轮训练。这对持续采集数据的产线场景特别重要。拿到数据后建议先跑一段脚本验证目录完整性别等训练报错了再回头查。下面这段遍历 train/val/test统计每个类别文件夹下的图片数量import os data_root micro_milling_wear for split in [train, val, test]: split_path os.path.join(data_root, split) if not os.path.isdir(split_path): print(f[ERROR] 缺少 {split} 目录) continue for class_name in sorted(os.listdir(split_path)): class_path os.path.join(split_path, class_name) count len([f for f in os.listdir(class_path) if f.lower().endswith((.jpg, .jpeg, .png))]) print(f{split}/{class_name}: {count} 张)这里我特意只统计 jpg/jpeg/png 结尾的文件防止把 macOS 的 .DS_Store、Windows 的 Thumbs.db 这类隐藏文件算进样本数。输出结果能暴露两类问题某个类别图片数为 0说明划分时漏了类某个类别数量明显偏少说明小样本问题比预期严重得提前决定是走数据增强还是类别加权。2.2 类别字典 JSON字段设计、读取和两种常见的映射方向类别字典文件是数据集的「翻译表」。图像分类模型输出的不是「轻微磨损」这种人类语言而是一个整数索引类别字典负责把整数索引翻译成类别名以及把文件夹名映射成整数索引。这两个方向相反缺一不可。实际项目里 JSON 有两种常见写法一种是索引到类别名形如{0: normal_wear}另一种是类别名到索引形如{normal_wear: 0}。这份数据集配套的类别字典用的是第一种典型结构长这样{ 0: normal_wear, 1: slight_wear, 2: moderate_wear, 3: severe_wear }读取并转换成训练和评估两边都需要的映射关系import json with open(category_dict.json, r, encodingutf-8) as f: id2label json.load(f) # 关键坑JSON 的 key 是字符串必须转成 int 才能当数值索引 id2label {int(k): v for k, v in id2label.items()} label2id {v: k for k, v in id2label.items()} print(索引 - 类别名:, id2label) print(类别名 - 索引:, label2id)两个细节容易翻车。第一json.load读出来的字典 key 全是字符串直接拿去做outputs.argmax()结果的索引会出现KeyError: 3这种诡异报错所以转完 dict 后必须把 key 转成 int。第二我同时维护 id2label 和 label2id 两个方向训练时用 label2id 把类别名转成标签评估和推理时用 id2label 把模型输出的索引转回类别名。两个方向都备好后面写训练脚本和画混淆矩阵就不用临时翻字典了。还有编码问题。Windows 上如果 JSON 是用记事本另存的默认可能是 GBK 编码open()里不写encodingutf-8会直接报UnicodeDecodeError如果 JSON 里带中文类别名读取后打印还会乱码。我的习惯是拿到文件先看编码统一转成 UTF-8 再进后续流程。提示拿到 JSON 后先跟 2.1 的目录扫描结果对一遍——JSON 里的类别名集合必须与 train 目录下的子文件夹名集合完全一致一个都不能多一个都不能少。2.3 划分比例与类别不均衡开训前的十分钟检查数据集已经划分好了但划分合理不等于类别均衡。微铣削刀具磨损天然不均衡正常磨损和轻微磨损的样本多严重磨损的样本少——刀具很少真磨到报废才换现场采集时严重磨损图像本来就难收集。这种不均衡会直接带偏分类器的决策边界让模型倾向于把模糊样本判到多数类。常规划分比例是 train:val:test 7:2:1 或 8:1:1这份数据具体怎么分的用脚本算一遍心里才有底def count_by_class(data_root): stats {} for split in [train, val, test]: split_path os.path.join(data_root, split) stats[split] {} for class_name in os.listdir(split_path): class_path os.path.join(split_path, class_name) stats[split][class_name] len(os.listdir(class_path)) return stats stats count_by_class(micro_milling_wear) for split, classes in stats.items(): total sum(classes.values()) print(f[{split}] 总数 {total}) for class_name, count in classes.items(): print(f {class_name}: {count} ({count / total:.1%}))如果发现某个类别占比很低两个标准补救方案一是类别重加权在损失函数里给少数类更高权重常用torch.nn.CrossEntropyLoss(weight...)权重按类别样本数的倒数归一化二是对少数类做更强的数据增强比如更高的旋转角度、额外的色彩抖动。我的习惯是先算占比再决定占比差距在两倍以内时数据增强往往比重加权更稳因为重加权本质上是放大少数类的梯度幅度太大会让整个训练震荡。3. 小样本条件下用 ResNet50 跑通全过程数据加载、训练参数与分类头处理3.1 从文件夹到 DataLoader一套直接可用的加载通路数据按文件夹保存那就直接用 ImageFolder 加载配合 transform 做预处理。微铣削磨损图像的典型问题是光照不均、刀面反光、拍摄角度有差异所以 transform 里除了 resize 和归一化我一般会加轻度色彩抖动。ResNet50 的输入尺寸是 224所有图像先统一到这个尺寸from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), 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]), ]) train_ds datasets.ImageFolder(micro_milling_wear/train, transformtrain_transform) val_ds datasets.ImageFolder(micro_milling_wear/val, transformval_transform) test_ds datasets.ImageFolder(micro_milling_wear/test, transformval_transform) train_loader DataLoader(train_ds, batch_size16, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size16, shuffleFalse, num_workers4)几个参数值得展开说。RandomHorizontalFlip对刀面图像是安全的刀具磨损程度不随水平翻转改变ColorJitter幅度我控制在 0.2太大会把本来就模糊的磨损纹理直接洗掉模型学不到真实特征。验证集和测试集只做 resize 和归一化不做随机增强——评估阶段任何随机性都会让指标不稳定这是拿到可复现结果的底线。batch_size16是针对小样本的保守选择。每类只有几十张图时batch 太大比如 64会造成一个 batch 里几乎全是同一类BatchNorm 层的均值和方差统计量剧烈抖动训练 loss 跟着震荡。验证集用shuffleFalse保证每个 batch 的构成稳定多次评估结果一致。3.2 小样本训练参数冻结、微调与数据增强的取舍小样本图像分类1-shot、5-shot 那种极限场景和常规分类的调参思路完全不同。微铣削磨损数据集每类可能只有几十到一两百张图直接从头训练 ResNet50 几乎必然过拟合——两千多万参数的模型容量几十张图根本喂不满。我的默认做法是用 ImageNet 预训练权重冻结浅层只微调后面的 stage 和分类头。浅层学到的是边缘、纹理、颜色块这类通用底层特征迁移到磨损图像上依然有效深层学的是与任务强相关的语义特征需要用少量样本去适配import torch import torchvision.models as models model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 冻结前两个 stage只训练 layer3、layer4 和分类头 for name, param in model.named_parameters(): if name.startswith(layer1) or name.startswith(layer2): param.requires_grad False # 分类头必须换ImageNet 是 1000 类磨损等级只有几类 num_classes len(train_ds.classes) model.fc torch.nn.Linear(model.fc.in_features, num_classes) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)layer1和layer2冻结后反向传播不会更新它们的参数显存占用和训练时间都有改善layer3、layer4保持可训练让高层特征向磨损纹理适配。分类头必须换掉in_features保留原值只改输出维度。训练参数我一般这样设参数建议值说明optimizerAdam, lr1e-4微调场景下 Adam 比 SGD 稳lr 比从头训练低一个量级schedulerCosineAnnealingLR小样本训练轮数少余弦退火比 StepLR 更平滑epochs30~50每类几十张图时30 轮足够再多就开始过拟合weight decay1e-4 ~ 5e-4正则化对抗小样本过拟合的关键这里有个容易翻车的点Adam 的 lr 用 1e-3 在这么小的数据量下偏大前几个 epoch loss 看着在降后面直接发散。我从 1e-4 起步loss 下降太慢再提到 5e-4千万不要反过来从大往小调。数据增强在小样本场景里是「后悔药」级别的存在。除了 ColorJitter强烈建议加 RandAugmentfrom torchvision.transforms import RandAugment train_transform.transforms.insert(2, RandAugment(num_ops2, magnitude9))num_ops2表示每次随机选两种变换magnitude9是强度。对磨损图像magnitude 超过 10 会让磨损区域的纹理失真模型学到的是增强后的假纹理而不是真实磨损特征这个值需要根据自己的图像微调。3.3 评估阶段分类头调整与 1-shot/5-shot 下的表现差异最近不少人问「用 ViT 评估时分类头用调整吗」换成 ResNet50 其实是同一个逻辑。答案分两种情况如果你加载的是 ImageNet 预训练权重直接做特征提取分类头还是 1000 类必须换掉如果你训练时已经替换过分了类头评估阶段直接调用即可不需要再动。评估代码里有几个必须注意的细节model.eval() # 必须放在循环外切换 BN/Dropout 行为 correct 0 total 0 all_preds [] all_labels [] with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) print(fVal Accuracy: {correct / total:.4f})model.eval()切换 BatchNorm 和 Dropout 的行为训练时 BN 用当前 batch 的统计量评估时用训练阶段积累的全局统计量。忘记这一行小 batch 下 BN 统计量漂移准确率掉三到五个点这是最常见的翻车点之一。torch.no_grad()关闭梯度计算省显存也提速评估阶段本来就不需要反向传播。关于 1-shot/5-shot这份数据集不是标准的 few-shot benchmark但可以拿来模拟小样本评估。常见做法是每类只取 1 张或 5 张训练其余全部做测试。做这种模拟时分类头也要处理把 fc 层重新随机初始化只保留 backbone 的预训练权重否则前一次实验的语义记忆会干扰新任务。每次 shot 实验都新建模型固定随机种子保证可复现。我实测的体感是ResNet50 在 5-shot 下能到八成左右准确率1-shot 直接掉到六成附近差距主要来自磨损等级边界模糊——轻微磨损和中度磨损的图像差异有时候人眼都难分辨。如果目标场景就是 1-shot建议换自监督预训练的 ViT 或 SimCLR 特征ResNet50 在极限小样本下的特征表达力确实不太够。4. 微铣削刀具磨损识别的避坑记录五个真实翻车现场4.1 坑一JSON 类别名和文件夹名不一致标签悄悄错位现象训练 loss 降得很顺利验证准确率也不低训练曲线一切正常。把模型部署到现场推理时却发现模型把「严重磨损」判成「正常磨损」而且错误高度集中在相邻类别上完全不像是随机犯错。原因category_dict.json 里的类别名顺序和 ImageFolder 扫描文件夹的字母序不一致。ImageFolder 按文件夹名字母序生成索引JSON 如果按「磨损程度从轻到重」的语义顺序写两边就对不上标签整体错位。解决用 2.2 的 label2id 做一次反向校验以文件夹名为准重建映射# train_ds 来自 3.1 的 ImageFolder folder_classes set(train_ds.classes) # ImageFolder 的类别顺序 json_classes set(id2label.values()) print(文件夹有但 JSON 没有:, folder_classes - json_classes) print(JSON 有但文件夹没有:, json_classes - folder_classes)差集为空后再开训这一步花两分钟能省后面两个小时的排查时间。4.2 坑二图里混入非图像文件DataLoader 加载到一半崩掉现象训练到中途报UnidentifiedImageError或者进度条卡在某个 batch 不动排查半天发现某个类别文件夹里混了一张扩展名是 .png 但内容不是图像的坏文件。原因ImageFolder 只按扩展名筛选文件不看文件内容。损坏文件被当成正常样本收进数据集解码时才报错。报错位置随机很难定位是哪个文件。解决训练前对全量做一次 PIL 解码验证把坏文件提前揪出来。下面的脚本沿用 2.1 的 data_rootfrom PIL import Image for split in [train, val, test]: for class_name in os.listdir(os.path.join(data_root, split)): class_dir os.path.join(data_root, split, class_name) for fname in os.listdir(class_dir): fpath os.path.join(class_dir, fname) try: with Image.open(fpath) as img: img.load() except Exception: print(f[损坏文件] {fpath})img.load()是关键——Image.open()是惰性的只读文件头load()才会真正把像素数据读进内存解码错误在这一步暴露。整个数据集跑一遍只要几分钟能避免后面几小时的排查。4.3 坑三小样本直接从零训练 ResNet50过拟合到怀疑人生现象训练集准确率冲到 99%验证集只有 70%而且 epoch 越多验证集准确率不升反降典型的过拟合曲线。原因模型容量远大于数据量。ResNet50 两千多万参数每类几十张图根本填不满模型把训练集的噪声、背景纹理、光照特征全记下来了。解决切到预训练权重加冻结前两层加强数据增强的方案也就是 3.2 的做法。如果这样还过拟合把 ResNet50 换成 ResNet18——参数少一个量级小样本下往往更稳最终准确率未必比 ResNet50 差。别迷信大模型数据量摆在那模型再大也得有样本喂。4.4 坑四验证时忘了 model.eval()同一批图两次结果不一样现象训练循环里验证准确率 85%训练完单独加载模型再跑同一批验证图只有 79%两次结果不一致第一反应是权重保存出了问题其实权重好着呢。原因没切model.eval()的话BatchNorm 持续用当前 batch 的均值和方差做归一化。验证 batch 小、类别分布不均时BN 统计量每次都不一样结果自然不可复现。解决验证和测试的循环前固定写model.eval()训练循环前写model.train()。推理部署同理加载权重后第一件事就是 eval。这个习惯养成后我再没被这种「玄学」坑过。4.5 坑五工业图像光照不均归一化参数别照抄 ImageNet现象模型收敛慢训练 loss 震荡最后准确率卡在及格线上不去换什么学习率都救不回来。原因微铣削刀面图像是工业相机拍的强反光和暗区阴影是常态像素分布和 ImageNet 自然图像差异很大。照抄 ImageNet 的 mean/std0.485, 0.456, 0.406不是不能用但收敛速度和稳定度都会变差。解决用数据集本身的统计量重算 mean 和 std花十分钟loader DataLoader(train_ds, batch_size32, shuffleFalse, num_workers4) channels_sum, channels_sq_sum, num_batches 0, 0, 0 for images, _ in loader: channels_sum images.mean(dim[0, 2, 3]) channels_sq_sum (images ** 2).mean(dim[0, 2, 3]) num_batches 1 mean channels_sum / num_batches std (channels_sq_sum / num_batches - mean ** 2) ** 0.5 print(均值:, mean) print(标准差:, std)注意此时 images 已经经过 ToTensor 归一化到 [0,1] 区间算出来的 mean/std 也是这个量级直接替换 3.1 的Normalize参数即可。这个做法对工业图像特别值我见过不止一个项目把这一行换掉后收敛速度和最终准确率同时改善。5. 进阶验证用混淆矩阵和 Grad-CAM 判断模型学到了什么准确率只是分数对刀具磨损这种安全敏感场景模型错在哪、依据什么判断比分数本身更重要。跑出第一个模型我做的第一件事不是调参而是画混淆矩阵行是真实类别列是预测类别对角线外就是错位分布。微铣削磨损的典型结果是轻微磨损和中度磨损互相误判因为这两类图像差异本来就是一个渐变过程没有硬边界。这时候别急着加数据先确认标注本身是不是有歧义如果混淆方向是严重磨损被误判成中度磨损说明模型在「偏保守」这比反过来安全得多。混淆矩阵看宏观错误Grad-CAM 看微观依据——它回答「模型依据图像的哪个区域做判断」。热力图集中在刀刃磨损带说明学到了真东西打在背景或夹具上说明模型在走捷径泛化能力堪忧。核心实现基于 torch 的 hook 机制不需要额外装包import cv2 import numpy as np def grad_cam(model, image_tensor, target_class, device): model.eval() image_tensor image_tensor.unsqueeze(0).to(device) activations {} def hook_fn(module, input, output): activations[feat] output handle model.layer4[-1].register_forward_hook(hook_fn) output model(image_tensor) model.zero_grad() output[0, target_class].backward() grad activations[feat].grad.squeeze().cpu().numpy() feat activations[feat].squeeze().cpu().numpy() weights grad.mean(axis(1, 2), keepdimsTrue) cam (weights * feat).sum(axis0) cam np.maximum(cam, 0) cam cv2.resize(cam, (224, 224)) handle.remove() return (cam - cam.min()) / (cam.max() - cam.min() 1e-8)forward hook 抓取 layer4 最后一个卷积块的输出用目标类别梯度对通道加权平均得到每个空间位置对决策的贡献度最后 resize 回输入尺寸叠加在原图上。hook 用完后必须remove()否则后续推理会重复累积。从那以后我每次拿到新的图像分类数据集都强制走一遍「目录扫描 → 类别字典校验 → 混淆矩阵 → Grad-CAM」这套流程。文件夹结构和 JSON 字典虽然不起眼却是整个训练流程的地基地基歪了模型再先进也白搭。希望这份拆解能帮你在微铣削刀具磨损识别上少走几步弯路。本文还有配套的精品资源点击获取 
长河君 · 2026/10/11 6:35:54 · 阅读 1.2k 微铣削刀具磨损图像识别：可直接训练的划分好数据集与ResNet50实战
简介面向图像分类与刀具磨损检测任务的微铣削刀具磨损图像数据集适合正在学习YOLOv5分类、CNN图像识别或需要可直接训练的数据集的学生与开发者使用。资源已按训练集、测试集划分共两类样本对应好、坏等磨损状态具体类别映射见JSON类别字典文件训练集3600张、测试集900张可配合YOLOv5分类工程或自建分类网络直接训练目录下按类别文件夹存放接入数据加载器时几乎不需要额外整理。包内共2000个文件以1998张JPG图片为主另含1个Python可视化脚本和1个JSON字典压缩包大小165.2MBshow脚本可随机展示图片方便快速核对标注与划分是否合理。已有502人学习下载适合需要快速获取规范划分数据、完成磨损状态识别实验或复现分类流程的入门及进阶用户。1. 微铣削刀具磨损图像识别一份划分好、能直接开训的图像分类数据集做微铣削加工的人都有这个痛点刀具磨损直接决定工件表面质量换刀不及时一批零件就废。主流做法是拍刀面图像交给图像分类模型判断磨损等级省去停机目测。但这个场景数据量小、类别边界模糊更麻烦的是很多数据集只给原始图片train/val/test 划分要自己动手类别字典还得对着文件夹名手搓。这份资源把这两件烦心事做完了磨损图像按类别分好、按文件夹保存训练集、验证集、测试集划分现成还带一份类别字典 JSON 文件把类别索引和磨损等级一一对应。用 torchvision 的 ImageFolder 可以直接加载省掉整理数据的功夫把精力放在模型选型和调参上。适合两类人做刀具状态监测、想快速跑出磨损识别基线的工程师拿小样本图像分类练手、需要干净标准数据集的算法同学。下面从数据集结构、JSON 字段、训练参数到避坑点完整拆一遍。2. 数据集与类别字典吃透文件夹结构和 JSON 字段再开训拿到资源先别急着训练。文件夹保存加类别字典 JSON 这套组合看着简单但里面有两处映射关系必须提前盘清楚文件夹名到整数索引的映射、整数索引到磨损等级名称的映射。搞反了或者没对齐训练再久结果都是错的。2.1 文件夹保存的目录结构为什么说它天生适配 ImageFolder这份数据集按「文件夹保存」组织目录层级是标准的分类数据集约定micro_milling_wear/ ├── train/ │ ├── normal_wear/ │ │ ├── 001.jpg │ │ └── ... │ ├── slight_wear/ │ ├── moderate_wear/ │ └── severe_wear/ ├── val/ │ └── ... └── test/ └── ...每个磨损等级一个子文件夹文件夹名就是类别名。torchvision 的datasets.ImageFolder天然按这个结构加载扫描根目录下所有子文件夹每个子文件夹自动成为一个类别生成(image_path, class_index)样本对。你唯一需要记住的是ImageFolder 按文件夹名的字母序来分配索引不是按文件夹里的排列顺序也不是按你脑子里「从轻到重」的语义顺序。这个设计对微铣削磨损场景很友好。磨损等级通常分三到五类每类一个文件夹直观、好维护。相比 CSV 标注或 JSON 标注文件夹方式最大的好处是增量加图不用改标注文件——现场新拍一批图像丢进对应文件夹就能参与下一轮训练。这对持续采集数据的产线场景特别重要。拿到数据后建议先跑一段脚本验证目录完整性别等训练报错了再回头查。下面这段遍历 train/val/test统计每个类别文件夹下的图片数量import os data_root micro_milling_wear for split in [train, val, test]: split_path os.path.join(data_root, split) if not os.path.isdir(split_path): print(f[ERROR] 缺少 {split} 目录) continue for class_name in sorted(os.listdir(split_path)): class_path os.path.join(split_path, class_name) count len([f for f in os.listdir(class_path) if f.lower().endswith((.jpg, .jpeg, .png))]) print(f{split}/{class_name}: {count} 张)这里我特意只统计 jpg/jpeg/png 结尾的文件防止把 macOS 的 .DS_Store、Windows 的 Thumbs.db 这类隐藏文件算进样本数。输出结果能暴露两类问题某个类别图片数为 0说明划分时漏了类某个类别数量明显偏少说明小样本问题比预期严重得提前决定是走数据增强还是类别加权。2.2 类别字典 JSON字段设计、读取和两种常见的映射方向类别字典文件是数据集的「翻译表」。图像分类模型输出的不是「轻微磨损」这种人类语言而是一个整数索引类别字典负责把整数索引翻译成类别名以及把文件夹名映射成整数索引。这两个方向相反缺一不可。实际项目里 JSON 有两种常见写法一种是索引到类别名形如{0: normal_wear}另一种是类别名到索引形如{normal_wear: 0}。这份数据集配套的类别字典用的是第一种典型结构长这样{ 0: normal_wear, 1: slight_wear, 2: moderate_wear, 3: severe_wear }读取并转换成训练和评估两边都需要的映射关系import json with open(category_dict.json, r, encodingutf-8) as f: id2label json.load(f) # 关键坑JSON 的 key 是字符串必须转成 int 才能当数值索引 id2label {int(k): v for k, v in id2label.items()} label2id {v: k for k, v in id2label.items()} print(索引 - 类别名:, id2label) print(类别名 - 索引:, label2id)两个细节容易翻车。第一json.load读出来的字典 key 全是字符串直接拿去做outputs.argmax()结果的索引会出现KeyError: 3这种诡异报错所以转完 dict 后必须把 key 转成 int。第二我同时维护 id2label 和 label2id 两个方向训练时用 label2id 把类别名转成标签评估和推理时用 id2label 把模型输出的索引转回类别名。两个方向都备好后面写训练脚本和画混淆矩阵就不用临时翻字典了。还有编码问题。Windows 上如果 JSON 是用记事本另存的默认可能是 GBK 编码open()里不写encodingutf-8会直接报UnicodeDecodeError如果 JSON 里带中文类别名读取后打印还会乱码。我的习惯是拿到文件先看编码统一转成 UTF-8 再进后续流程。提示拿到 JSON 后先跟 2.1 的目录扫描结果对一遍——JSON 里的类别名集合必须与 train 目录下的子文件夹名集合完全一致一个都不能多一个都不能少。2.3 划分比例与类别不均衡开训前的十分钟检查数据集已经划分好了但划分合理不等于类别均衡。微铣削刀具磨损天然不均衡正常磨损和轻微磨损的样本多严重磨损的样本少——刀具很少真磨到报废才换现场采集时严重磨损图像本来就难收集。这种不均衡会直接带偏分类器的决策边界让模型倾向于把模糊样本判到多数类。常规划分比例是 train:val:test 7:2:1 或 8:1:1这份数据具体怎么分的用脚本算一遍心里才有底def count_by_class(data_root): stats {} for split in [train, val, test]: split_path os.path.join(data_root, split) stats[split] {} for class_name in os.listdir(split_path): class_path os.path.join(split_path, class_name) stats[split][class_name] len(os.listdir(class_path)) return stats stats count_by_class(micro_milling_wear) for split, classes in stats.items(): total sum(classes.values()) print(f[{split}] 总数 {total}) for class_name, count in classes.items(): print(f {class_name}: {count} ({count / total:.1%}))如果发现某个类别占比很低两个标准补救方案一是类别重加权在损失函数里给少数类更高权重常用torch.nn.CrossEntropyLoss(weight...)权重按类别样本数的倒数归一化二是对少数类做更强的数据增强比如更高的旋转角度、额外的色彩抖动。我的习惯是先算占比再决定占比差距在两倍以内时数据增强往往比重加权更稳因为重加权本质上是放大少数类的梯度幅度太大会让整个训练震荡。3. 小样本条件下用 ResNet50 跑通全过程数据加载、训练参数与分类头处理3.1 从文件夹到 DataLoader一套直接可用的加载通路数据按文件夹保存那就直接用 ImageFolder 加载配合 transform 做预处理。微铣削磨损图像的典型问题是光照不均、刀面反光、拍摄角度有差异所以 transform 里除了 resize 和归一化我一般会加轻度色彩抖动。ResNet50 的输入尺寸是 224所有图像先统一到这个尺寸from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), 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]), ]) train_ds datasets.ImageFolder(micro_milling_wear/train, transformtrain_transform) val_ds datasets.ImageFolder(micro_milling_wear/val, transformval_transform) test_ds datasets.ImageFolder(micro_milling_wear/test, transformval_transform) train_loader DataLoader(train_ds, batch_size16, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size16, shuffleFalse, num_workers4)几个参数值得展开说。RandomHorizontalFlip对刀面图像是安全的刀具磨损程度不随水平翻转改变ColorJitter幅度我控制在 0.2太大会把本来就模糊的磨损纹理直接洗掉模型学不到真实特征。验证集和测试集只做 resize 和归一化不做随机增强——评估阶段任何随机性都会让指标不稳定这是拿到可复现结果的底线。batch_size16是针对小样本的保守选择。每类只有几十张图时batch 太大比如 64会造成一个 batch 里几乎全是同一类BatchNorm 层的均值和方差统计量剧烈抖动训练 loss 跟着震荡。验证集用shuffleFalse保证每个 batch 的构成稳定多次评估结果一致。3.2 小样本训练参数冻结、微调与数据增强的取舍小样本图像分类1-shot、5-shot 那种极限场景和常规分类的调参思路完全不同。微铣削磨损数据集每类可能只有几十到一两百张图直接从头训练 ResNet50 几乎必然过拟合——两千多万参数的模型容量几十张图根本喂不满。我的默认做法是用 ImageNet 预训练权重冻结浅层只微调后面的 stage 和分类头。浅层学到的是边缘、纹理、颜色块这类通用底层特征迁移到磨损图像上依然有效深层学的是与任务强相关的语义特征需要用少量样本去适配import torch import torchvision.models as models model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 冻结前两个 stage只训练 layer3、layer4 和分类头 for name, param in model.named_parameters(): if name.startswith(layer1) or name.startswith(layer2): param.requires_grad False # 分类头必须换ImageNet 是 1000 类磨损等级只有几类 num_classes len(train_ds.classes) model.fc torch.nn.Linear(model.fc.in_features, num_classes) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)layer1和layer2冻结后反向传播不会更新它们的参数显存占用和训练时间都有改善layer3、layer4保持可训练让高层特征向磨损纹理适配。分类头必须换掉in_features保留原值只改输出维度。训练参数我一般这样设参数建议值说明optimizerAdam, lr1e-4微调场景下 Adam 比 SGD 稳lr 比从头训练低一个量级schedulerCosineAnnealingLR小样本训练轮数少余弦退火比 StepLR 更平滑epochs30~50每类几十张图时30 轮足够再多就开始过拟合weight decay1e-4 ~ 5e-4正则化对抗小样本过拟合的关键这里有个容易翻车的点Adam 的 lr 用 1e-3 在这么小的数据量下偏大前几个 epoch loss 看着在降后面直接发散。我从 1e-4 起步loss 下降太慢再提到 5e-4千万不要反过来从大往小调。数据增强在小样本场景里是「后悔药」级别的存在。除了 ColorJitter强烈建议加 RandAugmentfrom torchvision.transforms import RandAugment train_transform.transforms.insert(2, RandAugment(num_ops2, magnitude9))num_ops2表示每次随机选两种变换magnitude9是强度。对磨损图像magnitude 超过 10 会让磨损区域的纹理失真模型学到的是增强后的假纹理而不是真实磨损特征这个值需要根据自己的图像微调。3.3 评估阶段分类头调整与 1-shot/5-shot 下的表现差异最近不少人问「用 ViT 评估时分类头用调整吗」换成 ResNet50 其实是同一个逻辑。答案分两种情况如果你加载的是 ImageNet 预训练权重直接做特征提取分类头还是 1000 类必须换掉如果你训练时已经替换过分了类头评估阶段直接调用即可不需要再动。评估代码里有几个必须注意的细节model.eval() # 必须放在循环外切换 BN/Dropout 行为 correct 0 total 0 all_preds [] all_labels [] with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) print(fVal Accuracy: {correct / total:.4f})model.eval()切换 BatchNorm 和 Dropout 的行为训练时 BN 用当前 batch 的统计量评估时用训练阶段积累的全局统计量。忘记这一行小 batch 下 BN 统计量漂移准确率掉三到五个点这是最常见的翻车点之一。torch.no_grad()关闭梯度计算省显存也提速评估阶段本来就不需要反向传播。关于 1-shot/5-shot这份数据集不是标准的 few-shot benchmark但可以拿来模拟小样本评估。常见做法是每类只取 1 张或 5 张训练其余全部做测试。做这种模拟时分类头也要处理把 fc 层重新随机初始化只保留 backbone 的预训练权重否则前一次实验的语义记忆会干扰新任务。每次 shot 实验都新建模型固定随机种子保证可复现。我实测的体感是ResNet50 在 5-shot 下能到八成左右准确率1-shot 直接掉到六成附近差距主要来自磨损等级边界模糊——轻微磨损和中度磨损的图像差异有时候人眼都难分辨。如果目标场景就是 1-shot建议换自监督预训练的 ViT 或 SimCLR 特征ResNet50 在极限小样本下的特征表达力确实不太够。4. 微铣削刀具磨损识别的避坑记录五个真实翻车现场4.1 坑一JSON 类别名和文件夹名不一致标签悄悄错位现象训练 loss 降得很顺利验证准确率也不低训练曲线一切正常。把模型部署到现场推理时却发现模型把「严重磨损」判成「正常磨损」而且错误高度集中在相邻类别上完全不像是随机犯错。原因category_dict.json 里的类别名顺序和 ImageFolder 扫描文件夹的字母序不一致。ImageFolder 按文件夹名字母序生成索引JSON 如果按「磨损程度从轻到重」的语义顺序写两边就对不上标签整体错位。解决用 2.2 的 label2id 做一次反向校验以文件夹名为准重建映射# train_ds 来自 3.1 的 ImageFolder folder_classes set(train_ds.classes) # ImageFolder 的类别顺序 json_classes set(id2label.values()) print(文件夹有但 JSON 没有:, folder_classes - json_classes) print(JSON 有但文件夹没有:, json_classes - folder_classes)差集为空后再开训这一步花两分钟能省后面两个小时的排查时间。4.2 坑二图里混入非图像文件DataLoader 加载到一半崩掉现象训练到中途报UnidentifiedImageError或者进度条卡在某个 batch 不动排查半天发现某个类别文件夹里混了一张扩展名是 .png 但内容不是图像的坏文件。原因ImageFolder 只按扩展名筛选文件不看文件内容。损坏文件被当成正常样本收进数据集解码时才报错。报错位置随机很难定位是哪个文件。解决训练前对全量做一次 PIL 解码验证把坏文件提前揪出来。下面的脚本沿用 2.1 的 data_rootfrom PIL import Image for split in [train, val, test]: for class_name in os.listdir(os.path.join(data_root, split)): class_dir os.path.join(data_root, split, class_name) for fname in os.listdir(class_dir): fpath os.path.join(class_dir, fname) try: with Image.open(fpath) as img: img.load() except Exception: print(f[损坏文件] {fpath})img.load()是关键——Image.open()是惰性的只读文件头load()才会真正把像素数据读进内存解码错误在这一步暴露。整个数据集跑一遍只要几分钟能避免后面几小时的排查。4.3 坑三小样本直接从零训练 ResNet50过拟合到怀疑人生现象训练集准确率冲到 99%验证集只有 70%而且 epoch 越多验证集准确率不升反降典型的过拟合曲线。原因模型容量远大于数据量。ResNet50 两千多万参数每类几十张图根本填不满模型把训练集的噪声、背景纹理、光照特征全记下来了。解决切到预训练权重加冻结前两层加强数据增强的方案也就是 3.2 的做法。如果这样还过拟合把 ResNet50 换成 ResNet18——参数少一个量级小样本下往往更稳最终准确率未必比 ResNet50 差。别迷信大模型数据量摆在那模型再大也得有样本喂。4.4 坑四验证时忘了 model.eval()同一批图两次结果不一样现象训练循环里验证准确率 85%训练完单独加载模型再跑同一批验证图只有 79%两次结果不一致第一反应是权重保存出了问题其实权重好着呢。原因没切model.eval()的话BatchNorm 持续用当前 batch 的均值和方差做归一化。验证 batch 小、类别分布不均时BN 统计量每次都不一样结果自然不可复现。解决验证和测试的循环前固定写model.eval()训练循环前写model.train()。推理部署同理加载权重后第一件事就是 eval。这个习惯养成后我再没被这种「玄学」坑过。4.5 坑五工业图像光照不均归一化参数别照抄 ImageNet现象模型收敛慢训练 loss 震荡最后准确率卡在及格线上不去换什么学习率都救不回来。原因微铣削刀面图像是工业相机拍的强反光和暗区阴影是常态像素分布和 ImageNet 自然图像差异很大。照抄 ImageNet 的 mean/std0.485, 0.456, 0.406不是不能用但收敛速度和稳定度都会变差。解决用数据集本身的统计量重算 mean 和 std花十分钟loader DataLoader(train_ds, batch_size32, shuffleFalse, num_workers4) channels_sum, channels_sq_sum, num_batches 0, 0, 0 for images, _ in loader: channels_sum images.mean(dim[0, 2, 3]) channels_sq_sum (images ** 2).mean(dim[0, 2, 3]) num_batches 1 mean channels_sum / num_batches std (channels_sq_sum / num_batches - mean ** 2) ** 0.5 print(均值:, mean) print(标准差:, std)注意此时 images 已经经过 ToTensor 归一化到 [0,1] 区间算出来的 mean/std 也是这个量级直接替换 3.1 的Normalize参数即可。这个做法对工业图像特别值我见过不止一个项目把这一行换掉后收敛速度和最终准确率同时改善。5. 进阶验证用混淆矩阵和 Grad-CAM 判断模型学到了什么准确率只是分数对刀具磨损这种安全敏感场景模型错在哪、依据什么判断比分数本身更重要。跑出第一个模型我做的第一件事不是调参而是画混淆矩阵行是真实类别列是预测类别对角线外就是错位分布。微铣削磨损的典型结果是轻微磨损和中度磨损互相误判因为这两类图像差异本来就是一个渐变过程没有硬边界。这时候别急着加数据先确认标注本身是不是有歧义如果混淆方向是严重磨损被误判成中度磨损说明模型在「偏保守」这比反过来安全得多。混淆矩阵看宏观错误Grad-CAM 看微观依据——它回答「模型依据图像的哪个区域做判断」。热力图集中在刀刃磨损带说明学到了真东西打在背景或夹具上说明模型在走捷径泛化能力堪忧。核心实现基于 torch 的 hook 机制不需要额外装包import cv2 import numpy as np def grad_cam(model, image_tensor, target_class, device): model.eval() image_tensor image_tensor.unsqueeze(0).to(device) activations {} def hook_fn(module, input, output): activations[feat] output handle model.layer4[-1].register_forward_hook(hook_fn) output model(image_tensor) model.zero_grad() output[0, target_class].backward() grad activations[feat].grad.squeeze().cpu().numpy() feat activations[feat].squeeze().cpu().numpy() weights grad.mean(axis(1, 2), keepdimsTrue) cam (weights * feat).sum(axis0) cam np.maximum(cam, 0) cam cv2.resize(cam, (224, 224)) handle.remove() return (cam - cam.min()) / (cam.max() - cam.min() 1e-8)forward hook 抓取 layer4 最后一个卷积块的输出用目标类别梯度对通道加权平均得到每个空间位置对决策的贡献度最后 resize 回输入尺寸叠加在原图上。hook 用完后必须remove()否则后续推理会重复累积。从那以后我每次拿到新的图像分类数据集都强制走一遍「目录扫描 → 类别字典校验 → 混淆矩阵 → Grad-CAM」这套流程。文件夹结构和 JSON 字典虽然不起眼却是整个训练流程的地基地基歪了模型再先进也白搭。希望这份拆解能帮你在微铣削刀具磨损识别上少走几步弯路。本文还有配套的精品资源点击获取