使用 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

用于加载图像的辅助函数（隐藏）

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

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

为了您的方便，将打印模型的句柄（网址）。那里提供了有关每个模型的更多文档。

选择图像分类模型

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

你可以选择以下图片之一，或使用你自己的图片。请记住，模型的输入大小各不相同，其中一些模型使用动态输入大小（支持对未缩放的图像进行推理）。鉴于此，load_image 方法会将图像重新调整为预期格式。

选择输入图像

2024-03-09 13:46:31.989759: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:282] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected

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.12 s, sys: 188 ms, total: 2.31 s
Wall time: 1.36 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 911 ms, sys: 103 ms, total: 1.01 s
Wall time: 80.1 ms
(1)   35 - blowing glass: 0.7747849822044373
(2)   34 - blowing bubble gum: 0.1064402163028717
(3)   37 - blowing nose: 0.00587468734011054
(4)  148 - drinking shots: 0.0025945263914763927
(5)   36 - blowing leaves: 0.002559840213507414

png

了解更多信息

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

如果你想查看更多图像模型，可以在 tfhub.dev 上查看。