工业缺陷检测第一版方案：基于图像分类的落地实践与避坑指南

简介面向工业质检与计算机视觉初学者提供一套基于图像分类的工业缺陷检测完整项目。资源围绕划痕、裂纹、异物等表面缺陷识别以视频讲解配合可运行代码覆盖项目介绍、数据集处理、自定义数据集构建、CNN模型定义、训练与预测全流程适合需要快速上手深度学习视觉检测的工程师或学生。包体共1823个文件压缩包约346.61MB其中mp4为六段操作讲解py为模型定义/训练/预测脚本bmp为工业样本图像xml与pth分别用于标注和模型权重目录结构清晰便于按模块学习。已有86人学习下载。通过该资源可获得从数据处理到模型部署的完整闭环代码与视频讲解逐段对应能够帮助读者理解缺陷分类任务中的关键参数与调优思路并复现可用的检测模型。1. 为什么工业缺陷检测的第一版方案我会先选图像分类人工质检工位面前是一张张快速流过的产品图像漏检的代价是整批退货。基于图像分类的工业缺陷检测就是要让模型先学会回答“这张图有没有问题、是哪类问题”输入一张产品图输出良品或缺陷类别附带一个置信度。很多人一上来就上目标检测我一般会劝住——工业项目首先追求的是数据成本可控、链路短、能快速上线图像分类只需要图像级标签不需要画框训练和部署都要简单得多。这篇笔记面向打算自己跑通这个方向的工程师从类别体系、数据增强、模型训练写到部署验证和避坑。你可以照着复现也能避开我在产线上踩过的那些坑。2. 图像分类在缺陷检测里的角色定位先解决“有没有问题”再谈定位2.1 先用分类还是先上目标检测三类现场场景的选型判断图像分类、目标检测、实例分割在缺陷检测里并不是“越往后越高级”的替代关系它们回答的是不同层级的问题。分类回答“这张图属于哪一类”检测在类别之外还想知道“缺陷在哪个位置、边界框多大”分割则要像素级轮廓。很多项目一启动就奔着目标检测理由是“以后反正要定位”——但边界框标注的成本是图像级标签的好几倍检测模型对数据量的胃口也更大在只有几百张缺陷图的阶段很容易过拟合。我判断选型只问三个问题。第一个问题缺陷在画面里的位置是不是相对固定固定机位加固定夹具的产线很常见比如注塑件端面、手机中框同一表面、螺丝端部这种场景缺陷永远落在相近区域分类模型完全能应对。第二个问题产线到底需要什么输出如果只需要自动分流为良品或不良品最多再按缺陷类型分拣分类就够只有返修工位需要知道缺陷坐标去打磨、补漆才必须让检测模型输出位置。第三个问题团队现在有多少已标注数据几百张图像级标签能开练几百张框级标注连热身都不够。我见过不少项目因为一上来就上YOLO标注搞了一个月还没进训练最后砍回分类才跑通。第一版先用分类打底等于把项目风险前置消掉了一部分。另外提一句最近一些项目尝试PatchCore这类基于特征嵌入的无监督异常检测适合缺陷样本非常稀缺的场景但它的阈值标定和特征提取器选择比较玄学换一个产品型号往往要重新过一遍正常样本。我当时复现PatchCore也折腾了几个晚上最后还是回到分类把基线先立住。图像分类的链路直观每类几十张样本就能出基线这也是我把它作为第一版方案的原因。2.2 类别体系怎么定从产线语料到 ImageFolder 目录拿到产线的缺陷代码表最容易犯的错是直接把“管理代码”当“训练标签”。产线的不良代码是为质量追溯设计的同一个代码下往往挂着多种外观完全不同的缺陷反过来两个不同代码的缺陷在人眼看来几乎一样。类别体系划分有一条底线同一类里的图像必须有一眼可辨的视觉共性不同类之间必须有稳定可区分的边界。我习惯按“返工动作”来划类别可以返修的划痕、压伤、只能报废的开裂、材料缺损、以及良品。返工动作决定了质检员接下来干什么也决定了缺陷在图像上的视觉分组是否合理。还有一类问题在工业场景里很隐蔽语义重叠。“脏污”“异物”“杂质”三个词在不同班组嘴里可能指的是同一种深色斑点如果按原始描述拆成三个类模型在像素层面找不到稳定差异训练时只会在这几类之间乱猜。宁可先合并成一个contam类后续样本多了再细分。类别体系一旦定下来改起来等于数据全部重来属于选了就难后悔的决定我在这上面栽过跟头。每类样本数也要同时盘点最少的那类如果不到30张先把容易混淆的类合并掉不要硬撑多分类。标签整理这一环节我习惯在训练脚本外面单独走一个脚本把产线导出的CSV映射成PyTorch ImageFolder兼容的目录结构脚本本身也很简单import csv import shutil from pathlib import Path # 产线导出的标签表表头至少包含文件名, 缺陷代码, 缺陷描述 # 示例行image_001.jpg, NG-SCRATCH, 表面划伤 LABEL_FILE defect_labels.csv RAW_DIR Path(raw_images) # 产线采集的原始图片目录 OUT_DIR Path(dataset/train) # 输出的 ImageFolder 训练目录 CATEGORY_MAP { # 业务代码 - 训练类别 OK: ok, NG-SCRATCH: scratch, NG-DENT: dent, NG-CONTAM: contam, # 视觉特征不清晰先合并 NG-FOREIGN: contam, } with open(LABEL_FILE, r, encodingutf-8) as f: reader csv.DictReader(f) for row in reader: code row[缺陷代码].strip() category CATEGORY_MAP.get(code) if category is None: continue # 未映射的代码先跳过不要硬分 src RAW_DIR / row[文件名] if not src.exists(): continue dst OUT_DIR / category / row[文件名] dst.parent.mkdir(parentsTrue, exist_okTrue) shutil.copy2(src, dst)这段脚本读CSV后按CATEGORY_MAP把业务缺陷代码映射成训练类别再复制进类别目录。CATEGORY_MAP是参数表按前面说的“视觉共性优先”原则调整脚本里NG-FOREIGN被并进contam就是一次主动合并。遇到未映射或文件不存在的行直接跳过产线表里经常混着“待定”“复判”这类非终态标签硬分进去会污染数据集。输出层只有四个目录模型任务也相应变成四分类第一版宁可少类也不要错类。验证集和测试集建议按生产批次单独整理不要在这个脚本里一起切避免同批次数据泄露到两个集合。跑完脚本后用几行代码盘点每类样本数这是开始训练前的第一道检查from torchvision import datasets train_ds datasets.ImageFolder(dataset/train) for cls_idx, cls_name in enumerate(train_ds.classes): count sum(1 for _, label in train_ds.samples if label cls_idx) print(f{cls_name:10s} {count})ImageFolder会按目录名自动生成标签目录名就是类别名所以子目录命名要用英文小写避免中文和空格带来的跨平台路径问题。如果最少类别的数量明显低于50张先别急着训练去现场补拍或者接受合并后的类别粒度。2.3 什么时候从分类升级到检测三个信号分类模型上线后不是直接替代人工我一般按“分流人工复核”的方式用模型输出softmax概率设双阈值p(ok)大于0.9的直接放行p(defect)大于0.8的直接拦截中间地带送人工复核。这套做法比追求单一准确率实在把人工注意力集中在最模糊的样本上。阈值不是拍脑袋定的先用验证集统计每个类别的置信度分布再取初值上线跑一周后微调。那什么时候需要从分类升级成检测模型三个信号。一是现场反馈需要知道缺陷坐标去返修比如打磨工位想知道划痕在哪个区域。二是同一张图里可能同时出现多种缺陷且处理动作不同分类只能给一个全局标签没法分别处理。三是改采集方案后视野变大缺陷在整张图里的占比变得很小这时分类模型容易被背景特征干扰检测的框相当于给模型一个强制注意力。出现其中之一再启动检测或分割方案才算合理的升级路径。在此之前分类模型的产出已经能把产线成本降下来一大块没必要为“高级”买单。3. 数据准备最容易翻车目录规范、读取与增强策略3.1 用 ImageFolder 管理缺陷图目录规范与读取代码工业缺陷数据有个特点图片数量大、类别分布极不均匀、且带有明显的批次特征。目录组织上我直接用PyTorch内置的ImageFolder格式它在torchvision里和datasets.ImageFolder直接对接跨机器拷贝时只要保持目录不变标签就跟着目录走比维护一份txt清单再对文件名靠谱得多。建立目录的脚本很简单cd /path/to/project mkdir -p dataset/train/{ok,scratch,dent,contam} mkdir -p dataset/val/{ok,scratch,dent,contam} mkdir -p dataset/test/{ok,scratch,dent,contam}train/val/test三套目录的划分有一个关键纪律按产线采集时间或生产批次切而不是随机打乱。同一个工件可能被拍了多张不同角度的照片随机划分会把同一工件的不同视角同时分进训练和验证验证集的准确率就会虚高上线后一换批次就原形毕露。我一般会把现场采集记录里的拍摄日期、线体编号作为划分依据例如前七天数据进训练第八天进验证再往后独立一天进测试。测试集必须是模型训练期间完全没见过的批次这样出来的指标才接近上线表现。目录建好后读取和盘点数据用datasets.ImageFolder训练前先跑一遍统计脚本from torchvision import datasets import collections train_ds datasets.ImageFolder(dataset/train) counter collections.Counter(label for _, label in train_ds.samples) for cls_idx, cls_name in enumerate(train_ds.classes): print(f{cls_name:10s} {counter[cls_idx]} 张) print(类别映射, train_ds.class_to_idx)ImageFolder从子目录名生成标签class_to_idx按目录名的ASCII码排序所以类别命名建议统一用小写英文。这个统计脚本能同时确认两件事目录结构是否完整以及每类样本数是否符合预期。样本数异常时训练脚本跑得越深浪费越多这里花三十秒就能拦下来。3.2 增强策略的取舍哪些变换对产线真实有效工业缺陷检测的增强策略和自然图像分类差别很大。自然图像可以用随机旋转、水平翻转因为目标不依赖方向但产线图像是固定机位加固定光源产品摆放方向一致很多缺陷具有方向性比如拉丝纹路的划痕只在特定角度出现。如果做了RandomRotation(90)一个竖向划痕会被转成横向模型被迫学习“旋转不变性”而真实产线根本不存在旋转这等于人为把分布域扩大反而降低判别力。我常用的增强组合只模拟产线真实会出现的波动光源老化带来的亮度下降、镜头轻微脏污带来的对比度变化、工件在夹具里的微小位置偏移。对应的transforms是这样写的from torchvision import transforms IMAGENET_MEAN [0.485, 0.456, 0.406] IMAGENET_STD [0.229, 0.224, 0.225] train_transforms transforms.Compose([ transforms.Resize((224, 224)), # 统一输入尺寸 transforms.RandomResizedCrop(224, scale(0.9, 1.0)), # 模拟夹具轻微位移 transforms.ColorJitter( brightness0.2, contrast0.2, saturation0.1 ), # 模拟光源波动 transforms.RandomAdjustSharpness(0.5, p0.3), # 模拟轻微失焦 transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), ]) val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), ])RandomResizedCrop的scale取(0.9, 1.0)只做轻微缩放裁剪模拟夹具位移带来的取景差异而不是随机裁剪到图像的一个角落。ColorJitter的brightness和contrast各给0.2对应光源老化和频闪波动saturation只给0.1因为工业图的颜色变化本身很小。RandomAdjustSharpness的sharpness_factor取0.5是模糊方向模拟镜头脏污和轻度失焦概率设为0.3防止过度模糊掩盖缺陷。验证集和测试集只用Resize加Normalize不做任何增强保证评估的是模型在真实采样分布上的表现。提示增强策略设完以后把增强后的图批量导出人眼过一遍。判断标准是——不借助原图一个人也能正确判断这张图属于哪个类别。如果增强让缺陷变得难以辨认它就不是增强而是噪声污染。缺陷样本量少且不均匀是常态光靠增强不够还要在采样层面做平衡。常见做法是用WeightedRandomSampler让样本少的类别在被采样时获得更高权重from torch.utils.data import WeightedRandomSampler import collections train_ds datasets.ImageFolder(dataset/train, transformtrain_transforms) counter collections.Counter(label for _, label in train_ds.samples) total sum(counter.values()) weights [total / counter[cls] for cls in range(len(train_ds.classes))] sample_weights [weights[label] for _, label in train_ds.samples] sampler WeightedRandomSampler( sample_weights, num_sampleslen(train_ds), replacementTrue )每个样本的采样权重等于“该类样本数的倒数”少类样本被抽中的概率相对提高。replacementTrue允许同一样本在一轮里重复出现相当于在数据层面模拟了过采样。num_samples默认取len(train_ds)会让每个epoch变长如果少数类特别少可以设成少数类样本数的三到五倍训练效率更高。使用sampler时DataLoader就不能再设shuffleTrue两者互斥。类别不均衡的另一个选择是直接在损失函数里加权我在下一章会用到两种方式也可以同时使用但要克制叠加过猛容易导致少类过拟合。4. 训练一个能过验收的基线模型模型选型与关键参数4.1 模型选型为什么先别追最新的图像分类模型工业缺陷检测项目选模型第一原则是“能复现、能部署、能解释”而不是追最新的图像分类模型刷榜单。最新的Transformer类模型在公开数据集上精度确实漂亮但工业图上可能只有几百张样本预训练权重来自自然图像领域差异大微调效果不一定比得上传统CNN而且部分结构的算子和部署工具链支持度差真正搬到工控机上才发现某个算子不支持就进退两难。我第一版固定用ResNet18理由很朴素结构简单、训练稳定、ONNX导出无坑。如果数据集超过一万张再对比EfficientNet-B0它的参数量更小推理更快但训练对学习率的敏感度稍高。产品表面纹理复杂、缺陷细碎的可以升到ResNet34但参数量增加带来的收益在二千张以内的小数据集上基本体现不出来。这个决策表是我常用的模型参数量训练稳定度合适的数据规模备注ResNet18约11M高几百到几千张第一版基线ResNet34约21M高上万张纹理细碎时再考虑EfficientNet-B0约5M中上万张部署资源紧时MobileNetV3约4M中几万张边缘盒子专用预训练权重建议用ImageNet预训练特征对边缘、纹理的底层响应是通用的工业缺陷一般也是从边缘纹理上区分。但工业图和自然图分布差异大微调学习率要放小见下面训练脚本里的参数。4.2 train.py 核心训练循环损失函数、学习率与早停训练脚本我用PyTorch写这里给出一个能直接跑通的完整核心循环包含预训练模型加载、类别加权损失、余弦退火学习率和验证集早停保存。数据量不大时建议先在本机CPU或单卡GPU上跑通再上更强的机器import torch import torch.nn as nn import collections from torch.utils.data import DataLoader, WeightedRandomSampler from torchvision import models, datasets, transforms def build_dataloaders(data_rootdataset, batch_size32): train_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomResizedCrop(224, scale(0.9, 1.0)), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.1), transforms.RandomAdjustSharpness(0.5, p0.3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) train_ds datasets.ImageFolder(data_root /train, transformtrain_tf) val_ds datasets.ImageFolder(data_root /val, transformval_tf) counter collections.Counter(label for _, label in train_ds.samples) total sum(counter.values()) cls_weights [total / counter[cls] for cls in range(len(train_ds.classes))] sample_weights [cls_weights[label] for _, label in train_ds.samples] sampler WeightedRandomSampler(sample_weights, num_sampleslen(train_ds), replacementTrue) train_loader DataLoader(train_ds, batch_sizebatch_size, samplersampler, num_workers4) val_loader DataLoader(val_ds, batch_sizebatch_size, shuffleFalse, num_workers4) return train_loader, val_loader, train_ds.classes device cuda if torch.cuda.is_available() else cpu num_classes 4 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) # 替换全连接层 model.to(device) criterion nn.CrossEntropyLoss( weighttorch.tensor(cls_weights, dtypetorch.float32).to(device) ) optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30) train_loader, val_loader, classes build_dataloaders() best_acc 0.0 for epoch in range(30): model.train() train_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(images), labels) loss.backward() optimizer.step() train_loss loss.item() * images.size(0) model.eval() correct total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) preds model(images).argmax(dim1) correct (preds labels).sum().item() total labels.size(0) val_acc correct / total if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fepoch {epoch:02d} val_acc{val_acc:.4f} 保存最优权重) scheduler.step()先说两个容易忽略的细节。model.fc的输入维度用model.fc.in_features动态读取不写死后面换EfficientNet时这行不用改。cls_weights在训练前已经算好类别少的类在CrossEntropyLoss里获得更高权重这是和WeightedRandomSampler并行的另一道保险。参数上lr3e-4是预训练权重微调的常规起点工业图和ImageNet分布差异大学到1e-3以上容易把预训练特征冲没loss先降后涨。weight_decay1e-4给AdamW加轻微正则小数据集上能明显抑制过拟合。CosineAnnealingLR的T_max30让学习率在30轮里从3e-4余弦下降到接近0省去手动调整的麻烦。早停用最朴素的方式验证集准确率只要创新高就保存权重训练结束后取best_model.pth即可。如果机器只有CPU把batch_size降到8到16ResNet18在数百张图上跑几十轮也就几十分钟能跑但别指望飞快。注意如果训练过程中出现loss为NaN优先检查学习率和数据里有没有损坏的图片文件不要先怀疑模型结构。把lr降到1e-4试跑20步是最快的排查手段。4.3 边训练边检查的三个点loss 曲线、混淆矩阵与 badcase 落盘只看最终准确率不够训练过程里我习惯盯三个东西。第一个是train loss和val_acc随epoch的变化train loss下降但val_acc停滞说明模型在过拟合此时权重衰减或增强强度再往上调一点两边都不动多半是学习率太小或类别本身不可分先去回看标注。第二个是验证集混淆矩阵混淆矩阵能直接指出哪些类别在模型眼里不可分。第三是badcase落盘把预测错的图按“真实类别_as_预测类别”的命名存下来训练完花十分钟人工过一遍比在指标上反复纠结有用得多。badcase落盘代码import torch import numpy as np from pathlib import Path from PIL import Image MEAN_NUMPY np.array([0.485, 0.456, 0.406], dtypenp.float32).reshape(1, 1, 3) STD_NUMPY np.array([0.229, 0.224, 0.225], dtypenp.float32).reshape(1, 1, 3) def dump_badcases(model, val_loader, class_names, save_dirbadcases, max_save50): model.eval() save_dir Path(save_dir) save_dir.mkdir(exist_okTrue) saved 0 with torch.no_grad(): for images, labels in val_loader: images images.to(device) preds model(images).argmax(dim1) for i, (pred, label) in enumerate(zip(preds, labels)): if pred label or saved max_save: continue img images[i].cpu().permute(1, 2, 0).numpy() img (img * STD_NUMPY MEAN_NUMPY) # 反标准化 img np.clip(img * 255, 0, 255).astype(uint8) img Image.fromarray(img) name f{class_names[label]}_as_{class_names[pred]}_{saved:03d}.jpg img.save(save_dir / name) saved 1 print(f已保存 {saved} 张 badcase 到 {save_dir})这段代码把验证集里预测错误的图片还原成可查看的JPEG。反标准化必须和训练时的Normalize一一对应MEAN和STD取同一个值否则导出的图颜色发灰发暗没法看。文件名里带上了真实类别和预测类别人工看的时候一眼就能判断是标注问题还是模型问题——比如大量的“contam_as_scratch”基本可以断定这两个类的视觉边界不清晰回到第2章做类别合并而不是继续调参。5. 避坑工业缺陷图像分类最常见的 5 个翻车现场5.1 验证集准确率 98%一上产线就崩数据划分泄漏现象本地训练验证集准确率接近98%badcase也少结果把模型接到产线当天不良品的漏检率比人工还高。原因训练验证集是按文件名随机划分的同一个工件的多张照片被拆到了两边模型等于提前看过了答案更隐蔽的是批次泄漏模具磨损是逐渐变化的前一批照片和后一批照片外观有系统性差异随机划分会让验证集里混进高度相似的样本。解决回到采集记录按拍摄时间或生产批次切分。前80%的批次进训练后20%的批次进验证测试集单独留一个完整批次。验证要模拟上线的场景——验证集在时间上必须晚于训练集指标才有参考价值。5.2 模型把“脏污”全学成“划痕”标注噪声比想象中大现象混淆矩阵里contam和scratch互相混人工翻badcase才发现不少训练图本身的标签就是错的同一批样本在不同人手里判出来不一样。原因缺陷标注本质是人工判断边缘case本来就界限模糊加上不同班组对缺陷代码的理解不一致标签噪声可能高达10%到20%。解决两步走。先训练一个基线模型把验证集上预测错误的图全部导出人工复核一遍改正明显的错标对复核时都说不清楚该归哪类的样本直接删掉或者合并类别。标签清洗一次到位很难要多轮迭代但每轮都能让模型稳定一点。5.3 环境搭建期卡住找不到 msvcp140.dll 导致代码跑不起来现象视频讲解和示例代码讲解里模型训练得很顺利自己在Windows电脑上按步骤装完环境一import torch就报“由于找不到msvcp140.dll无法继续执行代码”。原因PyTorch和OpenCV都是C扩展编译的Windows下依赖Visual C运行库纯Python环境装得再干净也绕过不了这一层。解决去微软官网搜索Visual C Redistributable把2015到2022的x64运行库装齐大多数情况马上解决如果还不行就用Anaconda新建一个干净环境通过conda安装PyTorch全家桶conda会把配套的运行库一起处理好。装完以后用一行代码验证环境import torch; print(torch.version)再往下跑训练脚本别等到训练中段才暴露环境问题。5.4 训练 loss 一开始就 NaN 或一直不降现象训练第一个epoch的loss就是nan或者loss曲线平平的一直在2.5上下浮动也就是多分类随机猜测的均值附近。原因学习率过大是NaN的最常见原因预训练模型微调给了1e-2会导致梯度爆炸数据里混入了损坏的图片文件比如0字节、截断的jpg也会让backward产出nan。而loss不降多半是预训练权重没加载成功、模型从头练起或者类别本身不可分。解决先把lr降到1e-4跑20步确认loss能下降再逐步调回写一个简单遍历DataLoader的脚本把打不开的图片文件打印出来删掉检查model.load_state_dict的输出确认没有missing keys的警告。这里每个检查只要几分钟却是定位问题的正道不要靠随机改参数碰运气。5.5 增强加了跟没加一样有时还更差现象在数据集上加了亮度和旋转增强验证集准确率不升反降badcase里还出现了训练时人为制造的“假样本”。原因工业场景本身光照恒定、方向固定增强如果引入了真实场景不存在的分布变化比如90度旋转、大幅色彩偏移模型学到的特征就会被带偏。解决先统计产线真实变化光源老化是渐变的亮度扰动范围给到±20%足够镜头脏污导致的是模糊而不是饱和度变化那就用模糊增强而不是色彩增强。每调完一次增强导出增强后的样本图人眼过一遍确认缺陷在人眼视角仍然清晰可辨再进入训练。6. 进阶用 Grad-CAM 验证模型看了哪里再用 ONNX 部署到产线6.1 Grad-CAM把黑匣子变成可解释的热力图分类模型在产线质检场景最大的问题是黑匣子——质量部门不会接受一个“只给结论不给理由”的系统。Grad-CAM能把模型决策依据可视化取最后一个卷积层的特征图用类别得分对特征图的梯度加权得到一张和原图同尺寸的热力图亮区就是模型判断时重点看的位置。实现上挂两个hook一个拿前向特征图一个拿反向梯度def hook_features(module, input, output): global feature_blob feature_blob output.detach() def hook_grad(module, grad_input, grad_output): global grad_blob grad_blob grad_output[0].detach() x x.to(device).requires_grad_(True) model.layer4[-1].register_forward_hook(hook_features) model.layer4[-1].register_full_backward_hook(hook_grad) def grad_cam(x, target_class): model.zero_grad() out model(x) out[0, target_class].backward() weight grad_blob.mean(dim(2, 3), keepdimTrue) cam (weight * feature_blob).sum(dim1, keepdimTrue) cam torch.relu(cam) return cam.squeeze().cpu().numpy()weight这里做的是全局平均池化把每个通道的梯度合成一个标量权重feature_blob是层4输出的特征图把通道加权求和再relu得到的就是热力图的原始响应。较新的PyTorch版本对register_backward_hook有弃用提示把函数名换成register_full_backward_hook即可。验证时挑预测错的图看热力图如果高亮区不在缺陷上而在背景或夹具上说明模型学的是批次或环境特征这比任何精度指标都更早暴露问题。6.2 ONNX 导出从 PyTorch 到产线工控机训练验证完了产线工控机一般没装PyTorch直接用ONNX Runtime推理最省事。导出代码很短model.eval() dummy torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy, defect_model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version13, )dynamic_axes把batch维度设为动态现场可以一次处理多张图不必每次重新导出。opset_version取13兼容性和算子支持度平衡得比较好。导出的onnx用onnxruntime跑一遍验证输入同样要经过Resize(224)和Normalize均值和方差和训练时保持一致这是上线前最容易漏的一步。我现在做新项目都默认先跑分类基线把badcase热力图发给工艺同事看他们一句话点出的问题往往比我在指标上抠两天更关键。希望帮到你。本文还有配套的精品资源点击获取