inception v4 官方实现_轻松学Pytorch迁移学习实现表面缺陷检查

最新推荐文章于 2023-11-20 17:05:32 发布

7323

最新推荐文章于 2023-11-20 17:05:32 发布

阅读量740

点赞数 1

文章标签： inception v4 官方实现

本文链接：https://blog.csdn.net/weixin_35822534/article/details/112197968

版权

本文介绍了如何使用Pytorch和torchvision进行迁移学习，以Inception V4模型为基础，对东北大学热轧带钢表面缺陷数据集进行图像分类。通过训练、测试和导出ONNX模型，实现OpenCV DNN调用部署。

摘要由CSDN通过智能技术生成

点击上方蓝字关注我们

微信公众号：OpenCV学堂

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

引言 ·

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

数据集

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

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

class SurfaceDefectDataset(Dataset):
    def __init__(self, root_dir):
        self.transform = transforms.Compose([transforms.ToTensor()])
        img_files = os.listdir(root_dir)
        self.defect_types =