PyTorch从零手写CNN实现NEU钢材表面缺陷检测（完整实战·单篇完结）

1. 项目概述

钢材表面缺陷检测是工业视觉检测的典型应用场景，在钢铁生产、智能制造、质量筛查领域具备极高实用价值。传统人工质检效率低、误差大、易受主观因素影响，基于深度学习的自动化缺陷识别可实现高精度、大批量、实时化检测。本文基于PyTorch从零搭建原生卷积神经网络，不使用预训练模型与迁移学习，完成NEU钢材表面六类缺陷分类任务。完整实现数据预处理、工业场景适配性数据增强、CNN模型搭建、训练调优、结果可视化、缺陷误判分析与工程优化方案，代码可直接运行、适配工业数据集特征，可用于工业视觉课题、课程设计、毕设落地项目。

2. 数据集介绍

NEU表面缺陷数据集（NEU Surface Defect Dataset）是钢材表面缺陷检测的标准公开数据集，专门用于金属表面缺陷分类、检测算法验证，数据贴合真实工业生产场景，包含光照不均、表面纹理干扰、缺陷尺度不一等真实工况问题。

数据集核心参数：

样本总量：1800张灰度钢材表面图像
类别数量：6类典型钢材缺陷
包含缺陷：轧制氧化皮、斑点、龟裂、划痕、压入氧化皮、夹杂
图像尺寸：统一200×200灰度图像
数据分布：每类缺陷各300张样本，类别均衡，无样本倾斜问题
任务属性：工业细粒度分类，同类缺陷形态差异大、不同类缺陷特征相似度高，检测难度高于通用物体分类

3. 环境依赖

项目适配Windows、Linux、Mac全平台，支持CPU与GPU训练，依赖通用深度学习基础库。

bash 复制代码

pip install torch torchvision matplotlib numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

4. 数据集预处理与数据增强

钢材缺陷图像为单通道灰度图，缺陷纹理细碎、边缘模糊、正负样本特征区分难度大。同时数据集样本总量较少，极易出现过拟合现象。因此本文针对性设计工业场景数据增强策略，在不破坏缺陷特征的前提下扩充样本多样性，提升模型泛化能力；测试集仅做标准化与尺寸归一化，保证评测结果真实可靠。

本次采用的增强策略：随机水平翻转、随机垂直翻转、小角度随机旋转、随机裁剪填充，适配工业缺陷无固定朝向的特征特点。

python 复制代码

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
import numpy as np

# 工业缺陷数据集归一化参数
mean = (0.5,)
std = (0.5,)

# 训练集增强：适配工业缺陷无固定方向特性
train_transform = transforms.Compose([
    transforms.Resize((200, 200)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(10),
    transforms.RandomCrop(200, padding=5),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

# 测试集仅标准化
test_transform = transforms.Compose([
    transforms.Resize((200, 200)),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

# 加载NEU钢材缺陷数据集
dataset = torchvision.datasets.ImageFolder(root="./NEU_defect", transform=train_transform)
# 划分训练集、测试集 8:2
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

# 替换测试集transform
test_dataset.dataset.transform = test_transform

# 分批加载数据
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=0)

print(f"训练集样本数：{len(train_dataset)}")
print(f"测试集样本数：{len(test_dataset)}")
print("缺陷类别：", dataset.class_to_idx)

# 样本可视化
def show_defect_sample():
    data_iter = iter(train_loader)
    images, labels = next(data_iter)
    img = torchvision.utils.make_grid(images[:6])
    img = img / 2 + 0.5
    np_img = img.numpy().transpose(1, 2, 0)
    plt.figure(figsize=(10,5))
    plt.imshow(np_img, cmap="gray")
    plt.title("钢材缺陷样本展示")
    plt.axis("off")
    plt.show()

show_defect_sample()

5. 自定义工业适配CNN网络模型

针对钢材灰度图像、细碎缺陷特征、小样本数据集特点，设计专用轻量化深度CNN网络。通过多层卷积逐级提取缺陷边缘、纹理、局部异常特征，配合池化层降维、多级Dropout正则化抑制小样本过拟合问题。网络结构深浅适中，兼顾细节特征提取能力与工业场景推理速度，适配缺陷细粒度分类任务。

python 复制代码

import torch.nn as nn
import torch.nn.functional as F

class NEU_Defect_CNN(nn.Module):
    def __init__(self, num_classes=6):
        super(NEU_Defect_CNN, self).__init__()
        # 第一层特征提取
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool1 = nn.MaxPool2d(2, 2)

        # 第二层特征提取
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(2, 2)

        # 第三层特征提取
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool3 = nn.MaxPool2d(2, 2)

        # 第四层特征提取
        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn4 = nn.BatchNorm2d(256)
        self.pool4 = nn.MaxPool2d(2, 2)

        # 正则化
        self.dropout = nn.Dropout(0.4)

        # 全连接分类层
        self.fc1 = nn.Linear(256 * 12 * 12, 1024)
        self.fc2 = nn.Linear(1024, num_classes)

    def forward(self, x):
        x = self.pool1(F.relu(self.bn1(self.conv1(x))))
        x = self.pool2(F.relu(self.bn2(self.conv2(x))))
        x = self.pool3(F.relu(self.bn3(self.conv3(x))))
        x = self.pool4(F.relu(self.bn4(self.conv4(x))))
        x = self.dropout(x)

        x = x.view(-1, 256 * 12 * 12)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 设备自动适配
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = NEU_Defect_CNN().to(device)
print("工业缺陷检测模型初始化完成")
print(model)

6. 模型训练与验证

针对小样本工业数据集，采用交叉熵损失函数适配多分类任务，Adam优化器实现自适应参数更新。训练阶段开启数据增强、BN层与Dropout正则化，提升泛化能力；测试阶段关闭随机操作，保证评估结果稳定。逐轮记录训练与测试损失、准确率，完整保留实验迭代数据。

python 复制代码

import torch.optim as optim

# 超参数设置
EPOCHS = 50
LR = 0.0008

# 损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)

# 指标保存列表
train_loss_list = []
train_acc_list = []
test_loss_list = []
test_acc_list = []

# 迭代训练
for epoch in range(EPOCHS):
    # 训练模式
    model.train()
    train_loss = 0.0
    train_correct = 0
    total_train = 0

    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()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total_train += labels.size(0)
        train_correct += (predicted == labels).sum().item()

    avg_train_loss = train_loss / len(train_loader)
    train_acc = 100 * train_correct / total_train

    # 测试模式
    model.eval()
    test_loss = 0.0
    test_correct = 0
    total_test = 0

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total_test += labels.size(0)
            test_correct += (predicted == labels).sum().item()

    avg_test_loss = test_loss / len(test_loader)
    test_acc = 100 * test_correct / total_test

    # 保存指标
    train_loss_list.append(avg_train_loss)
    train_acc_list.append(train_acc)
    test_loss_list.append(avg_test_loss)
    test_acc_list.append(test_acc)

    print(f"第{epoch+1:2d}轮 | 训练损失:{avg_train_loss:.4f} | 训练精度:{train_acc:.2f}% | 测试损失:{avg_test_loss:.4f} | 测试精度:{test_acc:.2f}%")

# 保存模型权重
torch.save(model.state_dict(), "neu_defect_cnn_best.pth")
print("工业缺陷检测模型权重保存成功")

7. 实验结果可视化

绘制模型训练全过程损失与准确率曲线，直观展示小样本数据集的收敛状态与拟合情况，精准判断过拟合、欠拟合问题。

python 复制代码

plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

plt.figure(figsize=(12, 5))
# 损失曲线
plt.subplot(1, 2, 1)
plt.plot(train_loss_list, label="训练损失", color="#e74c3c")
plt.plot(test_loss_list, label="测试损失", color="#3498db")
plt.title("损失变化曲线")
plt.xlabel("迭代轮数")
plt.ylabel("损失值")
plt.legend()

# 准确率曲线
plt.subplot(1, 2, 2)
plt.plot(train_acc_list, label="训练精度", color="#e74c3c")
plt.plot(test_acc_list, label="测试精度", color="#3498db")
plt.title("准确率变化曲线")
plt.xlabel("迭代轮数")
plt.ylabel("准确率(%)")
plt.legend()

plt.tight_layout()
plt.show()

8. 创新优化与工程拓展方向

引入注意力机制，强化缺陷区域特征权重，抑制钢材背景纹理干扰
添加学习率余弦退火衰减，优化模型后期收敛精度
引入混淆矩阵，精准分析各类缺陷的误判、漏判情况
替换轻量化骨干网络，适配工业设备端实时检测场景
增加缺陷定位可视化，实现分类+定位双重工业功能

9. 常见问题与解决方案

小样本严重过拟合：增大Dropout系数、加强数据增强、降低初始学习率
显存不足：减小batch_size，代码自动兼容CPU训练
缺陷分类精度低：延长迭代轮数、增加卷积层特征提取维度、添加BN层归一化
数据集读取失败：检查文件夹层级，保证数据集根目录下为类别子文件夹