本文展示了如何使用PyTorch实现基于VGG16的皮革瑕疵识别检测，介绍了数据预处理、模型训练及结果可视化，强调了模型优化和性能提升的多种方法。
基于卷神经网络VGG16的皮革瑕疵的识别检测（代码部分）
文章目录
一、前言
本次博客我们将利用卷神经网络VGG16进行有关皮革瑕疵的识别检测，所用到的深度学习框架为PyTorch，我是在谷歌机器学习云平台Colab上使用GPU进行训练的。
本期博客主要是展示代码部分，后续我将更新有关本次博客的论文理论部分，因为这是我做的一个《人工智能》课程的期末大作业。
数据集我才用的是kaggle官网中的Leather Defect detection and Classification这个数据集，大小为32MB，下面是摘取kaggle官网该数据集的介绍：
现代皮革工业专注于生产高品质皮革产品以保持市场竞争力。然而，在材料处理、鞣制和染色等制造过程的各个阶段都会引入各种皮革缺陷。皮革表面的人工检查具有主观性和不一致的性质；因此，机器视觉系统已被广泛用于皮革缺陷的自动检测。由于皮革局部区域的纹理图案模糊不清，性质微小，因此有必要开发合适的图像处理算法来定位皮革缺陷，如折叠痕迹、生长痕迹、脱粒、松粒和针孔。
本研究提出了基于深度学习神经网络的方法，使用机器视觉系统对皮革缺陷进行自动定位和分类。在这项工作中，流行的卷积神经网络使用不同皮革缺陷的皮革图像进行训练，并遵循类激活映射技术来定位皮革缺陷类的感兴趣区域。与该结果中比较的最先进的神经网络架构相比，发现诸如 Google net、Squeeze-net、RestNet 等卷积神经网络可提供更好的分类准确性。
kaggle官网中有关该数据集的实验都是基于tensorflow实现的，本次博客我在大佬的基础上将其修改成Pytorch版本并进行修改优化，因为时间原因我并没有优化的特别好，所以感兴趣的朋友可以再继续研究，需要该数据集和代码的也可也来找我。
二、准备工作
1、导入依赖项和数据集
import numpy as np
import pandas as pd
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from torchvision.models import vgg16
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
# 将模型移动到GPU上
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
datasets_dir = 'data/Leather Defect Classification'
2、数据预处理
# 创建图像预处理的转换函数
transform = transforms.Compose([
transforms.Resize((227, 227)),
transforms.Grayscale(num_output_channels=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# 创建训练数据集和验证数据集
train_data = datasets.ImageFolder(datasets_dir, transform=transform)
val_data = datasets.ImageFolder(datasets_dir, transform=transform)
# 创建创建训练数据集的数据加载器和验证数据集的数据加载器
train_loader = DataLoader(train_data, batch_size=64, shuffle=True, num_workers=4)
val_loader = DataLoader(val_data, batch_size=64, shuffle=False, num_workers=4)
# 获取并显示训练数据中的类名
class_names = train_data.classes
print(class_names)
['Folding marks', 'Grain off', 'Growth marks', 'loose grains', 'non defective', 'pinhole']
3、可视化训练集中的图像样本
plt.figure(figsize=(10, 10))
for images, labels in train_loader:
for i in range(9):
ax = plt.subplot(3, 3, i + 1)
image = images[i].numpy().transpose((1, 2, 0))
image = (image * 0.5) + 0.5 # 反归一化
plt.imshow(image)
plt.title(class_names[labels[i]])
plt.axis("off")
if i == 8:
break
break
plt.show()
三、创建VGG16网络模型
在这里我们直接从网上下载VGG16的模型并稍加修改，感兴趣的朋友可以自行搭建VGG16模型或者构建其他模型进行实验。
# 定义模型结构
vgg_model = vgg16(pretrained=True) # 加载预训练的VGG16模型
for param in vgg_model.parameters():
param.requires_grad = False # 冻结所有参数，不进行梯度更新
# 修改最后的分类层
num_features = vgg_model.classifier[6].in_features # 获取原始模型的全连接层输入特征数
features = list(vgg_model.classifier.children())[:-1] # 获取原始模型除了最后一层全连接层外的所有层
features.extend([nn.Linear(num_features, len(class_names))]) # 添加一个新的全连接层，输出类别数目为len(class_names)
vgg_model.classifier = nn.Sequential(*features) # 替换原始模型的分类层
vgg_model.to(device) # 将模型移动到GPU上进行计算
VGG(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(inplace=True)
(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): ReLU(inplace=True)
(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13): ReLU(inplace=True)
(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(inplace=True)
(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(inplace=True)
(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(inplace=True)
(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(inplace=True)
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): ReLU(inplace=True)
(26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(27): ReLU(inplace=True)
(28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(29): ReLU(inplace=True)
(30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096, bias=True)
(1): ReLU(inplace=True)
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=6, bias=True)
)
)
四、模型训练
1、定义损失函数和优化器
criterion = nn.CrossEntropyLoss() # 定义交叉熵损失函数，用于多分类任务
optimizer = optim.Adam(vgg_model.parameters(), lr=0.001) # 使用Adam优化器，学习率为0.001，优化模型的可学习参数
2、训练模型
num_epochs = 10
train_losses = [] # 用于存储训练集每个epoch的损失值
val_losses = [] # 用于存储验证集每个epoch的损失值
train_corrects = 0 # 记录训练集正确预测的样本数量
val_corrects = 0 # 记录验证集正确预测的样本数量
for epoch in range(num_epochs):
vgg_model.train() # 设置为训练模式
train_loss = 0.0
train_corrects = 0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad() # 梯度清零
outputs = vgg_model(inputs) # 前向传播
_, preds = torch.max(outputs, 1) # 获取预测结果
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
train_loss += loss.item() * inputs.size(0)
train_corrects += preds.eq(labels).sum().item()
epoch_loss = train_loss / len(train_data)
epoch_acc = train_corrects / len(train_data)
train_losses.append(epoch_loss)
print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}, Acc: {epoch_acc:.4f}')
# 验证模型
vgg_model.eval() # 设置为评估模式
val_loss = 0.0
val_corrects = 0
for inputs, labels in val_loader:
inputs, labels = inputs.to(device), labels.to(device)
with torch.no_grad(): # 不计算梯度
outputs = vgg_model(inputs) # 前向传播
_, preds = torch.max(outputs, 1) # 获取预测结果
loss = criterion(outputs, labels) # 计算损失
val_loss += loss.item() * inputs.size(0)
val_corrects += preds.eq(labels).sum().item()
val_epoch_loss = val_loss / len(val_data)
val_epoch_acc = val_corrects / len(val_data)
val_losses.append(val_epoch_loss)
print(f'Validation Loss: {val_epoch_loss:.4f}, Acc: {val_epoch_acc:.4f}')
print('Finished Training')
3、保存模型
torch.save(vgg_model.state_dict(), 'trained_model.pth')
五、结果可视化
# 可视化训练过程中的损失
plt.figure()
plt.plot(range(1, num_epochs+1), train_losses, label='Training Loss')
plt.plot(range(1, num_epochs+1), val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()
# 打印最终准确率
print(f'Final Accuracy: {val_epoch_acc:.4f}')
下面是我测试的不同num_epochs下的模型训练结果：
1、num_epochs=10
Final Accuracy: 0.7281
2、num_epochs=20
Final Accuracy: 0.7353
3、num_epochs=30
Final Accuracy: 0.7508
4、num_epochs=40
Final Accuracy: 0.7444
从上面测试结果可以看出num_epochs=30的时候效果最佳。由于时间关系我就测试了这四次，感兴趣有时间的朋友可以再增加试试。
六、最后我想说
这个结果肯定不是最好的，我们还可以通过各种方法将模型准确率再提高，例如：
- 数据增强：应用随机的图像变换操作，如旋转、平移、缩放、翻转等，扩增训练数据集。这有助于提升模型的泛化能力。
- 学习率调度：使用学习率调度策略，如学习率衰减、余弦退火等，来动态调整学习率，有助于模型更好地收敛到最优解。
- 模型集成：将多个训练好的模型进行集成，可以通过投票、平均预测等方式来得到更准确的结果。
- 模型微调：除了修改分类层外，可以选择解冻一部分预训练模型的层，允许这些层的参数进行微调，从而更好地适应特定的任务。
- 更深的模型或更多的参数：尝试使用更深的卷积神经网络模型，或者增加模型的宽度（增加通道数、隐藏单元数等），可以提供更强大的表示能力。
- 集成多个模型架构：尝试使用不同的模型架构，如ResNet、Inception等，将它们组合在一起，以获得更好的模型性能。
- 更大的训练集：通过收集更多的训练数据，可以提供更丰富的样本分布，有助于模型更好地学习数据的特征。
- 正则化技术：使用正则化技术，如L1正则化、L2正则化、Dropout等，可以避免模型过拟合，提高泛化能力。
- 参数调优：尝试不同的优化器、损失函数、超参数等组合，并进行参数调优，以找到更好的模型配置。
- 迁移学习：使用预训练模型在相关任务上的权重作为初始参数，然后在新任务上进行微调，可以加速训练过程并提高模型的准确率。
在这里我就不再过多的叙述，也希望有感兴趣的朋友可以在此基础上继续研究。
扫一扫
2081
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
