pytorch自带网络_轻松学Pytorch迁移学习实现表面缺陷检查

点击上方蓝字关注我们

微信公众号：OpenCV学堂

关注获取更多计算机视觉与深度学习知识

 引言 ·

大家好，我又又好久没有更新这个系列了，主要原因归根结底只有一个懒，所谓 一勤天下无难事，百思心中有良谋。以后还争取每周更新，这次隔了一周没有更新，对不起大家了。今天给大家更新的是如何基于torchvision自带的模型完成图像分类任务的迁移学习，前面我们已经完成了对对象检测任务的迁移学习，这里补上针对图像分类任务的迁移学习，官方的文档比较啰嗦，看了之后其实可操作性很低，特别是对于初学者，估计看了之后就发懵的那种。本人重新改写了一波，代码简洁易懂，然后把训练结果导出ONNX，使用OpenCV DNN调用部署，非常实用！废话不多说了，少吹水！

数据集

东北大学热轧带钢表面缺陷数据集，该数据集是东北大学的宋克臣等几位老师收集的，一共包含了三类数据。这里使用(NEU surface defect database)，数据集收集了夹杂、划痕、压入氧化皮、裂纹、麻点和斑块总计6种缺陷，每种缺陷300张，图像尺寸为200×200。部分示例如下：

基于该数据集，实现pytorch数据类，完成数据集的加载与预处理的代码如下：

class SurfaceDefectDataset(Dataset):

怎么下载该数据集，后台回复"NEU"关键字即可获取下载地址

模型使用

Pytorchvison支持多种图像分类模型，这里我们选择残差网络模型作为迁移学习的基础模型，对输出层(最后一层)改为六个类别，其它特征层选择在训练时候微调参数。常见的ResNet网络模型如下：

基于ResNet18完成网络模型修改，最终的模型实现代码如下：

class SurfaceDefectResNet(torch.nn.Module):

模型训练与测试

模型训练跟前面讲的一些图像分类模型训练方式并无不同，基于交叉熵损失，完成训练，每个批次4张图像或者8张图，训练15个epoch之后，保存模型。然后使用模型测试35张测试图像，发现有两张预测错误，其余均正确。训练模型的代码如下：

# 训练模型的次数

转为为ONNX模式，OpenCV DNN部署调用，代码如下：

"surface_defect_resnet18.onnx")

预测运行结果如下：

运行结果与pytorch调用模型运行结果保持一致。由于这个是一个专栏，很多代码在以前的文章中已经给出了，这里就没有重复贴代码！

福利：免费送30课时OpenCV4 C++入门课程

扫码即可开始观看学习

因依老宿发心初

半学修心半读书

 推荐阅读 

OpenCV4系统化学习路线图-视频版本！

Tensorflow + OpenCV4 安全帽检测模型训练与推理

汇总 | OpenCV DNN模块中支持的分类网络

OpenCV中支持的人脸检测方法整理与汇总

详解ENet | CPU可以实时的道路分割网络

从Pytorch 的ONNX到OpenVINO中IR中间层

OpenCV 基于Inception模型图像分类

OpenCV4.4 + YOLOv4 真的可以运行了…..