ResNet18异常检测妙用：工业质检案例，1小时快速复现-开发者社区

ResNet18异常检测妙用：工业质检案例，1小时快速复现

引言：当工业质检遇上ResNet18

在工厂流水线上，质检员每天需要检查成千上万的零件是否有划痕、裂纹或装配缺陷。这种重复性工作不仅容易疲劳，还可能出现漏检。而ResNet18这个轻量级AI模型，就像一位不知疲倦的"数字质检员"，能快速识别产品表面的异常情况。

ResNet18是残差神经网络(ResNet)家族中最轻便的成员，只有18层深度。它最初设计用于图像分类任务，但通过迁移学习技巧，我们可以让它专注于工业质检这个特定场景。想象一下教一个会识别猫狗的大学生转行做质检员——我们只需要用少量工业缺陷图片对它进行"职业培训"，就能让它掌握专业的质检技能。

本文将带你用公开的MVTec工业异常检测数据集，1小时内完成从环境搭建到效果验证的全流程。即使你是第一次接触深度学习，也能跟着步骤快速看到实际效果，为后续申请企业级GPU服务器提供可靠依据。

1. 环境准备：10分钟搞定基础配置

1.1 选择适合的云GPU环境

工业质检通常需要处理高分辨率图片，建议选择至少8GB显存的GPU环境。在CSDN算力平台可以直接选择预装PyTorch的镜像，省去环境配置时间：

# 推荐镜像配置 PyTorch 1.12 + CUDA 11.3 Python 3.8 Ubuntu 20.04

1.2 安装必要依赖

启动环境后，只需安装几个额外包就能开始工作：

pip install torchvision==0.13.0 pip install opencv-python pip install matplotlib

2. 数据准备：获取工业质检数据集

2.1 下载MVTec数据集

MVTec是业界常用的工业异常检测基准数据集，包含多种工业品类的正常和缺陷样本：

import os import wget # 创建数据目录 os.makedirs('./mvtec', exist_ok=True) # 下载数据集（以PCB板子类为例） dataset_url = "https://www.mvtec.com/company/research/datasets/mvtec_ad/download/mvtec_anomaly_detection.tar.xz" wget.download(dataset_url, out='./mvtec/mvtec.tar.xz') # 解压数据 os.system('tar -xf ./mvtec/mvtec.tar.xz -C ./mvtec')

2.2 数据集结构解析

解压后的目录结构如下：

mvtec/ └── pcb/ ├── train/ # 正常样本 │ └── good/ # 无缺陷图片 ├── test/ # 测试样本 │ ├── good/ # 正常样本 │ └── defective/ # 各种缺陷类型 └── ground_truth/ # 缺陷标注掩码

3. 模型构建：改造ResNet18成为缺陷侦探

3.1 加载预训练模型

我们基于PyTorch加载预训练的ResNet18，并改造其最后一层：

import torch import torch.nn as nn from torchvision import models # 加载预训练模型 model = models.resnet18(pretrained=True) # 冻结所有卷积层参数（只训练最后一层） for param in model.parameters(): param.requires_grad = False # 替换最后一层（原分类1000类改为2类：正常/缺陷） num_features = model.fc.in_features model.fc = nn.Linear(num_features, 2) # 转移到GPU device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device)

3.2 数据预处理管道

工业图片需要统一尺寸和标准化处理：

from torchvision import transforms # 训练数据增强 train_transform = transforms.Compose([ transforms.Resize(256), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 测试数据处理 test_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])

4. 训练与评估：1小时见证AI质检员诞生

4.1 快速训练脚本

使用迁移学习只需少量epoch就能见效：

from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder import torch.optim as optim # 准备数据集 train_dataset = ImageFolder('./mvtec/pcb/train', transform=train_transform) test_dataset = ImageFolder('./mvtec/pcb/test', transform=test_transform) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.fc.parameters(), lr=0.001) # 训练循环 for epoch in range(10): # 只需10个epoch model.train() for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 每个epoch后评估 model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Epoch {epoch+1}, 准确率: {100 * correct / total:.2f}%')

4.2 可视化检测结果

训练完成后，我们可以直观查看模型表现：

import matplotlib.pyplot as plt import numpy as np def imshow(inp, title=None): """显示张量图像""" inp = inp.numpy().transpose((1, 2, 0)) mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) inp = std * inp + mean inp = np.clip(inp, 0, 1) plt.imshow(inp) if title is not None: plt.title(title) plt.pause(0.001) # 获取一批测试图像 images, labels = next(iter(test_loader)) images, labels = images.to(device), labels.to(device) # 预测 outputs = model(images) _, preds = torch.max(outputs, 1) # 显示结果 plt.figure(figsize=(10, 10)) for i in range(min(9, images.size(0))): # 最多显示9张 plt.subplot(3, 3, i+1) imshow(images[i].cpu()) plt.title(f'预测: {"缺陷" if preds[i] else "正常"}\n实际: {"缺陷" if labels[i] else "正常"}') plt.show()

5. 优化技巧：让质检更精准的3个秘诀

5.1 注意力热图定位缺陷

通过类激活映射(CAM)可以直观看到模型关注的区域：

from torch.nn.functional import adaptive_avg_pool2d def show_cam(image, model): # 获取最后一层卷积特征 features = model.conv1(image.unsqueeze(0)) features = model.layer1(features) features = model.layer2(features) features = model.layer3(features) features = model.layer4(features) # 获取分类权重 weights = model.fc.weight[1] # 缺陷类的权重 # 计算CAM cam = (weights.view(1, -1, 1, 1) * features).sum(1) cam = cam.squeeze().cpu().numpy() cam = np.maximum(cam, 0) # ReLU cam = cv2.resize(cam, (224, 224)) cam = cam - np.min(cam) cam = cam / np.max(cam) # 叠加到原图 image = image.cpu().numpy().transpose(1, 2, 0) image = np.array(255*(image - np.min(image))/np.ptp(image)).astype(np.uint8) heatmap = cv2.applyColorMap(np.uint8(255*cam), cv2.COLORMAP_JET) result = heatmap * 0.3 + image * 0.7 plt.imshow(result) plt.title('缺陷区域热力图') plt.axis('off') plt.show() # 对测试图像生成热图 sample_image = images[0] show_cam(sample_image, model)

5.2 处理类别不平衡问题

工业数据中正常样本通常远多于缺陷样本，可以采用加权损失：

# 计算类别权重 normal_count = len([x for x in os.listdir('./mvtec/pcb/train/good') if x.endswith('.png')]) defect_count = len([x for x in os.listdir('./mvtec/pcb/test/defective') if x.endswith('.png')]) class_weights = torch.tensor([1.0, normal_count/defect_count]).to(device) # 使用加权损失 criterion = nn.CrossEntropyLoss(weight=class_weights)

5.3 模型轻量化部署技巧

实际部署时可以考虑模型量化减小体积：

# 量化模型 quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 ) # 保存量化模型 torch.save(quantized_model.state_dict(), 'quantized_resnet18.pth')