TensorFlow 2之迁移学习分类 CIFAR-10

  |   0 评论   |   0 浏览

背景

使用卷积神经网络来分类 CIFAR-10 只能达到68%的正确率。

本文使用使用 TFHub 中的预训练模型 ImageNet 进行迁移学习,来实现更好的图像分类。

简介

ImageNet

TFHub 上有很多预训练好的模型(pretrained model),这次我们选择 ImageNet。ImageNet 数据集大约有1500万张图片,2.2万类,可以说你能想到,想象不到的图片都能在里面找到。想下载感受一下的话可以到官网下载ImageNet

当然每次训练不太可能使用所有的图片,一般使用子集,比如2012年ILSVRC分类数据集使用了大概1/10的图片。我们今天用于迁移学习的预训练模型就只有1001个分类,想知道这1001类分别有哪些可以看这里

初体验

备注:下面代码使用内存较多(50G+),本机无法运行。

# 迁移学习
#   在实际的应用中,预训练好的模型的输入输出可能并不能满足我们的需求,
#   另外,训练上百万甚至上千万张图片,可能需要花费好几天的时间,
#   那有没有办法只使用训练好的模型的一部分呢?
#   训练好的模型的前几层对特征提取有非常好的效果,如果可以直接使用,那就事半功倍了。
#   这种方法被称之为迁移学习(transfer learning)。

#   在接下来的例子中,我们复用了 ImageNet Classifier 的特征提取的部分,并定义了自己的输出层。
#   因为原来的模型输出是1001个分类,而我们希望识别的 CIFAR-10 数据集只有10个分类。

# 第一步:resize数据集
#   ImageNet Classifier 的输入固定为(224, 224, 3),但 CIFAR-10 数据集中的图片大小是 32 * 32,
#   简单起见,我们将每一张图片大小从 32x32 转换为 224x224,使用pillow库提供的 resize 方法。

import numpy as np
from PIL import Image
import matplotlib.pylab as plt
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras import layers, datasets


def resize(d, size=(224, 224)):
    return np.array([np.array(Image.fromarray(v).resize(size, Image.ANTIALIAS))
                     for i, v in enumerate(d)])


data_count = 50000
(train_x, train_y), (test_x, test_y) = datasets.cifar10.load_data()
train_x, test_x = resize(train_x[:data_count])/255.0, resize(test_x)/255.0
train_y = train_y[:data_count]

# 第二步:下载特征提取层
#   TFHub 提供了 ImageNet Classifier 去掉了最后的分类层的版本,可以直接下载使用。
#
# feature_extractor_url = 'https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4'
# 这里手动下载解压至 datasets/imagenet_feature_vector 目录中,离线加载。
feature_extractor_url = "datasets/imagenet_feature_vector"
feature_extractor_layer = hub.KerasLayer(feature_extractor_url,
                                         input_shape=(224, 224, 3))
# 这一层的训练值保持不变
feature_extractor_layer.trainable = False

# 第三步:添加分类层
#   我们在特征提取层后面,添加了输出为10的全连接层,用于最后的分类。
#   从model.summary()中我们可以看到特征提取层的输出是1280。
model = tf.keras.Sequential([
    feature_extractor_layer,
    layers.Dense(10, activation='softmax')
])
model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss='sparse_categorical_crossentropy',
              metrics=['acc'])
# model.summary()
# Model: "sequential"
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #
# =================================================================
# keras_layer (KerasLayer)     (None, 1280)              2257984
# _________________________________________________________________
# dense (Dense)                (None, 10)                12810
# =================================================================
# Total params: 2,270,794
# Trainable params: 12,810
# Non-trainable params: 2,257,984
# _________________________________________________________________

# 第四步:训练并评估模型
#   本文的示例模型非常简单,在feature_extractor_layer直接添加了输出层,可训练参数很少。
#   而且只使用大约一半的训练集,正确率仍然达到了83% 。
history = model.fit(train_x, train_y, epochs=1)
loss, acc = model.evaluate(test_x, test_y)
print(acc)
# 0.8379



结果

32/30000 [..............................] - ETA: 22:16 - loss: 2.8715 - acc: 0.09382020-10-03 00:21:51.635404: W tensorflow/core/framework/allocator.cc:107] Allocation of 154140672 exceeds 10% of system memory.
   64/30000 [..............................] - ETA: 12:41 - loss: 2.8005 - acc: 0.12502020-10-03 00:21:51.827406: W tensorflow/core/framework/allocator.cc:107] Allocation of 154140672 exceeds 10% of system memory.
   96/30000 [..............................] - ETA: 9:25 - loss: 2.8067 - acc: 0.1042 2020-10-03 00:21:52.025631: W tensorflow/core/framework/allocator.cc:107] Allocation of 154140672 exceeds 10% of system memory.
  128/30000 [..............................] - ETA: 7:50 - loss: 2.7670 - acc: 0.11722020-10-03 00:21:52.222147: W tensorflow/core/framework/allocator.cc:107] Allocation of 154140672 exceeds 10% of system memory.
30000/30000 [==============================] - 182s 6ms/sample - loss: 0.6055 - acc: 0.7961
10000/10000 [==============================] - 62s 6ms/sample - loss: 0.4705 - acc: 0.8379
0.8379

参考

附录

tensorflow-hub安装

pip3 install tensorflow-hub