TensorFlow2学习十五、使用VGG16模型训练自己的数据集

一、说明

VGG16在2014年ImageNet比赛中获胜。ImageNet数据集中有1000个图像属于1000个不同的类别。
VGG模型的权重是免费的，可以在您自己的模型和应用程序中加载和使用。这使得其他研究人员和开发人员可以在自己的工作和程序中使用最先进的图像分类模型。

二、实现过程

1. 这里使用谷歌的花卉数据集。

from __future__ import absolute_import, division, print_function
from tqdm import tqdm
from numpy.random import randn
 
import pathlib
import random
import matplotlib.pyplot as plt
 
import tensorflow as tf
import numpy as np
 
from matplotlib.image import imread
from keras.preprocessing import image
 
tf.enable_eager_execution()
 
AUTOTUNE = tf.data.experimental.AUTOTUNE
 
data_dir = tf.keras.utils.get_file('flower_photos','https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz', untar=True)
data_dir = pathlib.Path(data_dir)
 
label_names={'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}
label_key=['daisy','dandelion','roses','sunflowers','tulips']

2. 拆分数据集

all_images = list(data_dir.glob('*/*'))
all_images = [str(path) for path in all_images]
random.shuffle(all_images)
 
all_labels=[label_names[pathlib.Path(path).parent.name] for path in all_images]
 
data_size=len(all_images)
 
train_test_split=(int)(data_size*0.2)
 
x_train=all_images[train_test_split:]
x_test=all_images[:train_test_split]
 
y_train=all_labels[train_test_split:]
y_test=all_labels[:train_test_split]
 
IMG_SIZE=160
 
BATCH_SIZE = 32
 
def _parse_data(x,y):
  image = tf.read_file(x)
  image = tf.image.decode_jpeg(image, channels=3)
  image = tf.cast(image, tf.float32)
  image = (image/127.5) - 1
  image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
 
  return image,y
 
def _input_fn(x,y):
  ds=tf.data.Dataset.from_tensor_slices((x,y))
  ds=ds.map(_parse_data)
  ds=ds.shuffle(buffer_size=data_size)
  
  
  ds = ds.repeat()
  
  ds = ds.batch(BATCH_SIZE)
  
  ds = ds.prefetch(buffer_size=AUTOTUNE)
  
  return ds
  
train_ds=_input_fn(x_train,y_train)
validation_ds=_input_fn(x_test,y_test)

3. 从VGG16创建迁移学习模型

我们将从VGG16模型创建一个基本模型。这是在ImageNet数据集（一个包含140万张图像和1000类web图像的大型数据集）上预先训练的。

最后一个分类层不是很有用。取而代之的是，我们将遵循通常的做法，转而依赖于展平操作之前的最后一层。这个层被称为“瓶颈层”。与最终/顶层相比，瓶颈特性保留了许多通用性。

首先，实例化一个预先加载了ImageNet上训练的权重的VGG16模型。通过指定include_top=False参数，可以加载不包含分类层的网络。

IMG_SHAPE = (IMG_SIZE, IMG_SIZE, 3)
VGG16_MODEL=tf.keras.applications.VGG16(input_shape=IMG_SHAPE,
                                               include_top=False,
                                               weights='imagenet')
VGG16_MODEL.trainable=False
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
prediction_layer = tf.keras.layers.Dense(len(label_names),activation='softmax')
model = tf.keras.Sequential([
  VGG16_MODEL,
  global_average_layer,
  prediction_layer
])

4. 编译模型

model.compile(optimizer=tf.train.AdamOptimizer(), 
              loss=tf.keras.losses.sparse_categorical_crossentropy,
              metrics=["accuracy"])

5. 训练

history = model.fit(train_ds,
                    epochs=100, 
                    steps_per_epoch=2,
                    validation_steps=2,
                    validation_data=validation_ds)

6. 评估模型

validation_steps = 20
 
loss0,accuracy0 = model.evaluate(validation_ds, steps = validation_steps)
 
print("loss: {:.2f}".format(loss0))
print("accuracy: {:.2f}".format(accuracy0))

7. 打印学习曲线

plt.plot(history.history['acc'])
plt.plot(history.history['val_acc'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

# 损失函数
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()