使用 TensorFlow Hub 进行图像分类

在此 colab 中，您将尝试 TensorFlow Hub 中的多个图像分类模型，并确定哪一个最适合您的用例。

由于 TF Hub 鼓励对图像进行操作的模型采用一致的输入约定，因此很容易试验不同的架构来找到最适合您需求的架构。

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

import requests
from PIL import Image
from io import BytesIO

import matplotlib.pyplot as plt
import numpy as np

2022-12-14 22:22:00.284948: 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 22:22:00.285044: 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 22:22:00.285055: 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.

Helper functions for loading image (hidden)

选择一个图像分类模型。之后，设置一些内部变量，下载标签文件并准备使用。

模型之间存在一些技术差异，例如不同的输入大小、模型大小、准确性和推断时间。您可以在此处更改要使用的模型，直到找到最适合您用例的模型。

为方便起见，系统会打印模型的句柄 (url)。还提供了有关每个模型的更多文档。

注：所有这些模型都是在 ImageNet 数据集上训练的

Select an Image Classification model

Selected model: efficientnetv2-s : https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_s/classification/2
Images will be converted to 384x384

您可以选择以下图像之一，也可以使用您自己的图像。只需记住，模型的输入大小各不相同，其中一些使用动态输入大小（实现对未缩放图像的推断）。鉴于此，方法 load_image 已经将图像重新缩放为预期格式。

Select an Input Image

png

现在已选择模型，使用 TensorFlow Hub 加载它很简单。

这还会调用具有随机输入的模型作为“热身”运行。随后的调用通常要快得多，您可以将其与下面的延迟进行比较。

注：使用动态大小的模型可能需要针对每个图像大小都运行“热身”。

classifier = hub.load(model_handle)

input_shape = image.shape
warmup_input = tf.random.uniform(input_shape, 0, 1.0)
%time warmup_logits = classifier(warmup_input).numpy()

CPU times: user 2.2 s, sys: 509 ms, total: 2.7 s
Wall time: 2.74 s

一切准备就绪，可以进行推断了。这里您可以看到模型针对选定图像得出的排名前 5 的结果。

# Run model on image
%time probabilities = tf.nn.softmax(classifier(image)).numpy()

top_5 = tf.argsort(probabilities, axis=-1, direction="DESCENDING")[0][:5].numpy()
np_classes = np.array(classes)

# Some models include an additional 'background' class in the predictions, so
# we must account for this when reading the class labels.
includes_background_class = probabilities.shape[1] == 1001

for i, item in enumerate(top_5):
  class_index = item if includes_background_class else item + 1
  line = f'({i+1}) {class_index:4} - {classes[class_index]}: {probabilities[0][top_5][i]}'
  print(line)

show_image(image, '')

CPU times: user 20.3 ms, sys: 97 µs, total: 20.4 ms
Wall time: 19.8 ms
(1)   35 - blowing glass: 0.774784505367279
(2)   34 - blowing bubble gum: 0.10644055157899857
(3)   37 - blowing nose: 0.005874688737094402
(4)  148 - drinking shots: 0.002594532212242484
(5)   36 - blowing leaves: 0.0025598385836929083

png

了解更多

如果您想要了解更多信息并尝试使用这些模型进行迁移学习，您可以尝试此教程：图像分类的迁移学习

如果要查看更多图像模型，可以在 tfhub.dev 上检索