重新训练图像分类器

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 中查看源代码 下载笔记本 查看 TF Hub 模型

简介

图像分类模型有数百万个参数。从头训练需要大量带标签的训练数据和强大的算力。迁移学习会在新模型中重复使用已在相关任务上训练过的模型的一部分，可以显著降低这些需求。

此 Colab 演示了对于图像特征提取，如何使用 TensorFlow Hub 中的预训练 TF2 SavedModel（已在更庞大和更通用的 ImageNet 数据集上进行过训练），为五种花的分类构建 Keras 模型。您可以选择性地将特征提取器与新添加的分类器一起训练（“微调”）。

想要寻找替代工具？

这是 TensorFlow 编码教程。如果您只是希望找一款为其构建 TensorFlow 或 TF Lite 模型的工具，请查看通过 PIP 软件包 tensorflow-hub[make_image_classifier] 安装的 make_image_classifier 命令行工具，或查看此 TF Lite Colab。

设置

import itertools
import os

import matplotlib.pylab as plt
import numpy as np

import tensorflow as tf
import tensorflow_hub as hub

print("TF version:", tf.__version__)
print("Hub version:", hub.__version__)
print("GPU is", "available" if tf.config.list_physical_devices('GPU') else "NOT AVAILABLE")

2022-12-14 21:15:57.688879: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 21:15:57.688988: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 21:15:57.688999: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
TF version: 2.11.0
Hub version: 0.12.0
GPU is available

选择要使用的 TF2 SavedModel 模块

对于初学者，请使用 https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4。在代码中使用同一网址可以识别 SavedModel，在浏览器中访问该网址可以显示其文档。（请注意，TF1 Hub 格式的模型在此处无效。）

您可以在此处找到更多生成图像特征向量的 TF2 模型。

您可以尝试多个模型。您只需在以下单元上选择不同的模型，然后按照笔记本操作即可。

Selected model: efficientnetv2-xl-21k : https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet21k_xl/feature_vector/2
Input size (512, 512)

设置花数据集

输入根据所选模块适当调整大小。数据集扩充（即每次读取图像时的随机畸变）可改善训练，特别是在微调时。

data_dir = tf.keras.utils.get_file(
    'flower_photos',
    'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
    untar=True)

Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
228813984/228813984 [==============================] - 1s 0us/step


Found 3670 files belonging to 5 classes.
Using 2936 files for training.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
Found 3670 files belonging to 5 classes.
Using 734 files for validation.

定义模型

您需要做的就是使用 Hub 模块将一个线性分类器放在 feature_extractor_layer 的顶部。

为了提高速度，我们先从不可训练的 feature_extractor_layer 开始，但是，为了提高准确率，您也可以启用微调。

do_fine_tuning = False

print("Building model with", model_handle)
model = tf.keras.Sequential([
    # Explicitly define the input shape so the model can be properly
    # loaded by the TFLiteConverter
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(model_handle, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(len(class_names),
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))
model.summary()

Building model with https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet21k_xl/feature_vector/2
WARNING:tensorflow:Please fix your imports. Module tensorflow.python.training.tracking.data_structures has been moved to tensorflow.python.trackable.data_structures. The old module will be deleted in version 2.11.
Model: "sequential_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 keras_layer (KerasLayer)    (None, 1280)              207615832 
                                                                 
 dropout (Dropout)           (None, 1280)              0         
                                                                 
 dense (Dense)               (None, 5)                 6405      
                                                                 
=================================================================
Total params: 207,622,237
Trainable params: 6,405
Non-trainable params: 207,615,832
_________________________________________________________________

训练模型

model.compile(
  optimizer=tf.keras.optimizers.SGD(learning_rate=0.005, momentum=0.9), 
  loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1),
  metrics=['accuracy'])

steps_per_epoch = train_size // BATCH_SIZE
validation_steps = valid_size // BATCH_SIZE
hist = model.fit(
    train_ds,
    epochs=5, steps_per_epoch=steps_per_epoch,
    validation_data=val_ds,
    validation_steps=validation_steps).history

Epoch 1/5
183/183 [==============================] - 184s 838ms/step - loss: 0.9491 - accuracy: 0.8911 - val_loss: 0.7422 - val_accuracy: 0.9278
Epoch 2/5
183/183 [==============================] - 149s 816ms/step - loss: 0.6983 - accuracy: 0.9380 - val_loss: 0.6317 - val_accuracy: 0.9514
Epoch 3/5
183/183 [==============================] - 149s 815ms/step - loss: 0.5803 - accuracy: 0.9685 - val_loss: 0.5908 - val_accuracy: 0.9639
Epoch 4/5
183/183 [==============================] - 149s 815ms/step - loss: 0.5525 - accuracy: 0.9716 - val_loss: 0.5188 - val_accuracy: 0.9597
Epoch 5/5
183/183 [==============================] - 149s 815ms/step - loss: 0.5323 - accuracy: 0.9771 - val_loss: 0.5352 - val_accuracy: 0.9667

plt.figure()
plt.ylabel("Loss (training and validation)")
plt.xlabel("Training Steps")
plt.ylim([0,2])
plt.plot(hist["loss"])
plt.plot(hist["val_loss"])

plt.figure()
plt.ylabel("Accuracy (training and validation)")
plt.xlabel("Training Steps")
plt.ylim([0,1])
plt.plot(hist["accuracy"])
plt.plot(hist["val_accuracy"])

[<matplotlib.lines.Line2D at 0x7f49f4f03490>]

png

png

在验证数据中的图像上尝试模型：

x, y = next(iter(val_ds))
image = x[0, :, :, :]
true_index = np.argmax(y[0])
plt.imshow(image)
plt.axis('off')
plt.show()

# Expand the validation image to (1, 224, 224, 3) before predicting the label
prediction_scores = model.predict(np.expand_dims(image, axis=0))
predicted_index = np.argmax(prediction_scores)
print("True label: " + class_names[true_index])
print("Predicted label: " + class_names[predicted_index])

png

1/1 [==============================] - 5s 5s/step
True label: sunflowers
Predicted label: sunflowers

最后，您可以保存训练的模型，以便部署到 TF Serving 或 TF Lite（在移动设备上），如下所示。

saved_model_path = f"/tmp/saved_flowers_model_{model_name}"
tf.saved_model.save(model, saved_model_path)

WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 1594). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: /tmp/saved_flowers_model_efficientnetv2-xl-21k/assets
INFO:tensorflow:Assets written to: /tmp/saved_flowers_model_efficientnetv2-xl-21k/assets

可选：部署到 TensorFlow Lite

TensorFlow Lite 可让您将 TensorFlow 模型部署到移动和 IoT 设备上。下面的代码演示了如何将训练的模型转化为 TF Lite 以及应用 TensorFlow Model Optimization Toolkit 中的训练后工具。最后，它会在 TF Lite Interpreter 中运行模型，以检查结果质量。

  • 未使用优化的转换结果与之前相同（由舍入误差决定）。
  • 进行了优化但未使用任何数据的转换会将模型权重量化为 8 位，但是，神经网络激活的推理仍使用浮点计算。这会将模型大小减小为约四分之一，同时可以改善移动设备上的 CPU 延迟。
  • 最重要的是，如果提供一个小参考数据集来校准量化范围，神经网络激活的计算也可以量化为 8 位整数。在移动设备上，这会进一步加快推断速度，并使其可以在 EdgeTPU 之类的加速器上运行。

Optimization settings

2022-12-14 21:31:01.185665: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:362] Ignored output_format.
2022-12-14 21:31:01.185713: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:365] Ignored drop_control_dependency.
Wrote TFLite model of 826217852 bytes.

interpreter = tf.lite.Interpreter(model_content=lite_model_content)
# This little helper wraps the TFLite Interpreter as a numpy-to-numpy function.
def lite_model(images):
  interpreter.allocate_tensors()
  interpreter.set_tensor(interpreter.get_input_details()[0]['index'], images)
  interpreter.invoke()
  return interpreter.get_tensor(interpreter.get_output_details()[0]['index'])

num_eval_examples = 50 
eval_dataset = ((image, label)  # TFLite expects batch size 1.
                for batch in train_ds
                for (image, label) in zip(*batch))
count = 0
count_lite_tf_agree = 0
count_lite_correct = 0
for image, label in eval_dataset:
  probs_lite = lite_model(image[None, ...])[0]
  probs_tf = model(image[None, ...]).numpy()[0]
  y_lite = np.argmax(probs_lite)
  y_tf = np.argmax(probs_tf)
  y_true = np.argmax(label)
  count +=1
  if y_lite == y_tf: count_lite_tf_agree += 1
  if y_lite == y_true: count_lite_correct += 1
  if count >= num_eval_examples: break
print("TFLite model agrees with original model on %d of %d examples (%g%%)." %
      (count_lite_tf_agree, count, 100.0 * count_lite_tf_agree / count))
print("TFLite model is accurate on %d of %d examples (%g%%)." %
      (count_lite_correct, count, 100.0 * count_lite_correct / count))

INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
TFLite model agrees with original model on 50 of 50 examples (100%).
TFLite model is accurate on 50 of 50 examples (100%).