4. 基于深度学习的钢材表面缺陷分类

任务目标

构建一个基于预训练CNN（ResNet18）的钢材表面缺陷分类器，对NEU-DET数据集中6类常见缺陷（划痕、麻点、夹杂、氧化皮、裂纹、轧制氧化坑）进行图像分类。最终交付可运行的训练与推理脚本，并给出验证精度的具体方法。

环境与依赖

• Python 3.8+，PyTorch 1.10+，torchvision 0.11+
• 数据集：NEU-DET（东北大学钢材表面缺陷数据集），包含1800张灰度图像（300×300），每类300张
• 目录结构：data/ 下按 train/ 与 val/ 划分，每个子目录内按类别存放图像

关键完整代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
import os

# ---------- 配置 ----------
DATA_ROOT = "data"          # 含 train/ 与 val/ 子目录
BATCH_SIZE = 32
EPOCHS = 10
LR = 1e-4
NUM_CLASSES = 6
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------- 数据预处理 ----------
train_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=3),  # 灰度图转3通道以适配ImageNet预训练
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=3),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# ---------- 数据加载 ----------
train_dataset = datasets.ImageFolder(os.path.join(DATA_ROOT, "train"), train_transform)
val_dataset = datasets.ImageFolder(os.path.join(DATA_ROOT, "val"), val_transform)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

# ---------- 模型构建 ----------
def build_model(num_classes):
    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
    # 冻结前几层，只微调最后两个stage和全连接层
    for name, param in model.named_parameters():
        if "layer4" not in name and "fc" not in name:
            param.requires_grad = False
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    return model

model = build_model(NUM_CLASSES).to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)

# ---------- 训练 ----------
def train_one_epoch():
    model.train()
    total_loss, correct, total = 0, 0, 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()
        total_loss += loss.item() * images.size(0)
        _, preds = torch.max(outputs, 1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

# ---------- 验证 ----------
def validate():
    model.eval()
    correct, total = 0, 0
    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)
    return correct / total

# ---------- 主流程 ----------
if __name__ == "__main__":
    best_acc = 0.0
    for epoch in range(EPOCHS):
        train_loss, train_acc = train_one_epoch()
        val_acc = validate()
        print(f"Epoch {epoch+1}/{EPOCHS} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}")
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), "steel_defect_model.pth")
    print(f"Best validation accuracy: {best_acc:.4f}")

运行方式

1. 下载NEU-DET数据集，按类别整理为 data/train/<类别名>/ 与 data/val/<类别名>/ 结构（建议按8:2划分训练/验证集）
2. 安装依赖：pip install torch torchvision
3. 执行训练：python train.py，模型权重保存为 steel_defect_model.pth

验证方法与判定标准

训练完成后观察终端输出的验证集准确率。判定通过的标准：

• 验证集准确率 ≥ 85%（在NEU-DET标准划分下，ResNet18微调通常可达90%以上）
• 训练过程中验证损失不再下降且准确率趋于平稳，说明模型已收敛
• 使用 torch.load 加载权重后，对单张缺陷图像执行前向推理，输出类别索引与对应缺陷名称一致

推理示例

def predict(image_path, model, class_names):
    from PIL import Image
    img = Image.open(image_path).convert("L")
    img = val_transform(img).unsqueeze(0).to(DEVICE)
    model.eval()
    with torch.no_grad():
        output = model(img)
        _, pred = torch.max(output, 1)
    return class_names[pred.item()]

# class_names 按 ImageFolder 中类别字母序排列
class_names = ["crazing", "inclusion", "patches", "pitted_surface", "rolled_in_scale", "scratches"]
print(predict("test_scratches.jpg", model, class_names))  # 预期输出: scratches