使用VGG模型自定义图像分类任务

ゝ一纸荒年。 2022-05-24 12:42 229阅读 0赞

# 前言 #

网上关于VGG模型的文章有很多，有介绍算法本身的，也有代码实现，但是很多代码只给出了模型的结构实现，并不包含数据准备的部分，这让人很难愉快的将代码迁移自己的任务中。为此，这篇博客接下来围绕着如何使用VGG实现自己的图像分类任务，从数据准备到实验验证。代码基于Python与TensorFlow实现，模型结构采用VGG-16，并且将很少的出现算法和理论相关的东西。

# 数据准备 #

## 下载数据和转换代码 ##

大多数人自己的训练数据，一般都是传统的图片形式，如.jpg，.png等等，而图像分类任务的话，这些图片的天然组织形式就是一个类别放在一个文件夹里，那么有啥大众化的数据集是这样的组织形式呢？TensorFlow的FlowersData，它下载下来是这个样子：  
![这里写图片描述][70]  
一共有五类，每一类中都有几百张图，我们把这些数据组织成TFrecord形式，对应的博客[在这里][Link 1]，源码的github[在这里][Link 2]，FlowersData数据集[在这里][Link 3]。  
有上面这三个东西之后，就可以生成TFrecord文件了。

## 组织图片数据 ##

首先将FlowersData文件夹下的数据分成两个部分，训练数据和测试数据，我把原文件五个类别中都拿出大概100张图左右，数据的构成和路径如下：  
![这里写图片描述][70 1]

## 生成训练TFrecord ##

#图片路径
    cwd = 'F:\\flowersdata\\trainimages\\'

#文件路径
    filepath = 'F:\\flowersdata\\tfrecord\\train\\'

classes=['daisy',
             'dandelion',
             'roses',
             'sunflowers',
             'tulips']

#tfrecords格式文件名
    ftrecordfilename = ("traindata.tfrecords-%.3d" % recordfilenum)

## 生成效果：   ![这里写图片描述][70 2]   生成预测TFrecord ##

#图片路径
    cwd = 'F:\\flowersdata\\testimages\\'

#文件路径
    filepath = 'F:\\flowersdata\\tfrecord\\test\\'

classes=['daisy',
             'dandelion',
             'roses',
             'sunflowers',
             'tulips']

#tfrecords格式文件名
    ftrecordfilename = ("testdata.tfrecords-%.3d" % recordfilenum)

生成效果：  
![这里写图片描述][70 3]

# 训练模型 #

## 初始权重与源码下载 ##

VGG-16的初始权重我上传到了百度云，[在这里][Link 4]下载；  
VGG-16源码我上传到了github，[在这里][Link 5]下载；

在源码中：  
train\_and\_val.py文件是最终要执行的文件，它定了训练和预测的过程；  
input\_data.py是将上一步中生成的TFRecord文件组织成batch的过程；  
VGG.py定义了VGG-16的网络结构；  
tool.py是最底层，定义了一些卷积池化等操作。

## 训练模型 ##

train\_and\_val.py文件修改：

if __name__=="__main__":
        train()
        #evaluate()

根据自己的路径修改：

#初始权重路径
    pre_trained_weights = 'vgg16_pretrain/vgg16.npy'
    #训练数据路径
    train_data_dir = 'F:\\flowersdata\\tfrecord\\train\\traindata.tfrecords*'
        test_data_dir = 
    #预测数据路径
    'F:\\flowersdata\\tfrecord\\test\\testdata.tfrecords*'
    #训练生成文件路径
    train_log_dir = 'logs/train/'
    #预测生成文件路径
    val_log_dir = 'logs/val/'

根据自己的显存容量修改：

IMG_W = 224
    IMG_H = 224
    BATCH_SIZE = 8

训练过程每50个step打印loss；  
每200个step计算一个batch中的准确率；  
每1000个step保存一次权重。

# 预测 #

train\_and\_val.py文件修改：

if __name__=="__main__":
        #train()
        evaluate()

#训练过程中生成的权重
    log_dir = 'logs/train/'
    #预测数据集路径
    test_data_dir = 'F:\\flowersdata\\tfrecord\\test\\testdata.tfrecords*'
    #用于生成tf文件的图片数量
    n_test = 502

打印测试样本总数；  
打印正确预测的样本总数；  
打印top\_1。

[70]: /images/20220524/19faac261882405c8eecb8b2389584ca.png
[Link 1]: https://blog.csdn.net/chaipp0607/article/details/72960028
[Link 2]: https://github.com/chaipangpang/Convert-between-TFRecords-and-Images
[Link 3]: https://pan.baidu.com/s/1XH8MYimaLN3Vu9NgII_hnA
[70 1]: /images/20220524/edf884e75f3f48e4bbfcaa12bac48dd1.png
[70 2]: /images/20220524/79894cf21f2540978b143ec31d7b6566.png
[70 3]: /images/20220524/4688b041e7994bebb81fff02e4550482.png
[Link 4]: https://pan.baidu.com/s/1UfHxzHwHChXNr0WXzSlQNg
[Link 5]: https://github.com/chaipangpang/VGG_tensorflow