pytorch+VGG16 双分支输入图像分类

原创

已于 2023-12-22 14:48:31 修改 · 8.3k 阅读

129 ·

CC 4.0 BY-SA版权

文章标签：

#神经网络 #深度学习 #pytorch

于 2021-04-15 14:59:11 首次发布

本文介绍了如何将VGG16神经网络修改为接受两个输入通道，分别处理原始和处理过的图像进行图像分类。作者通过修改网络结构，创建自己的数据集类，并调整训练过程，实现了双输入的图像分类模型。提供的代码示例展示了从数据预处理到训练的完整流程，并分享了训练过程中的效果。

将神经网络改为双输入通道 pytorch+VGG16 图像分类

23.12.22更新
背景介绍
解决方法
解决效果

23.12.22更新

陆续有一些朋友发邮件来问，所以我花时间整理了一个新版本，代码位于：https://download.csdn.net/download/qq_41638200/88652320?spm=1001.2014.3001.5503
可自行下载。在运行时，需要把train.py的第27、28行数据路径修改为自己所需的数据路径：

origin_image_root = '/groupshare_1/cat_dog_dataset/origin/'
filter_image_root = '/groupshare_1/cat_dog_dataset/filter/'

如果想要此代码来练手，需要训练数据集，可以发邮件联系我。邮箱地址在评论区。
-------------------------------------------------------------以下是原博文--------------------------------------------------------------------------------------

背景介绍

本人是个菜鸟，初学神经网络，第一个任务就是图像分类，使用的是VGG16网络。网络实现参考大佬博客：https://blog.csdn.net/qq_37541097/article/details/103482003
遇到了这么一个需求：假设有原始图像a，处理过的图像b，我想让图像a和b成对在神经网络中训练，即双分支输入，我该怎么修改代码呢？
经过一番折腾，我终于搞定了，写出来跟大家分享一下。

解决方法

首先明确一点，VGG神经网络的结构：卷积层先提取特征，然后全连接层做分类器。基于这个结构特性，我们可以在网络里放两个卷积层提取特征的结构，然后把两个特征拼在一起

解决步骤一修改神经网络

修改初始化函数，加上第二个提取特征的卷积层，之所以加上第二个，不共用之前的，是因为一般两个图像特征提取器参数不共享效果会好一点。但是这也不一定，毕竟神经网络炼丹嘛，有时候说不定共享参数反而更好。

def __init__(self, features, features1, num_classes=1000, init_weights=False):
        super(VGG, self).__init__()
        self.features = features  # 卷积层提取特征1 
        self.features1 = features1  # 卷积层提取特征2 
        self.classifier = nn.Sequential(  # 全连接层进行分类
            nn.Dropout(p=0.5),
            nn.Linear(1024 * 7 * 7, 2048),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(True),
            nn.Linear(2048, num_classes)
        )
        if init_weights:
            self._initialize_weights(