KerasHub 入门指南
作者： Matthew Watson, Jonathan Bischof
创建日期 2022/12/15
最后修改日期 2024/10/17
描述： KerasHub API 简介。
KerasHub 是一个预训练模型库，旨在做到简单、灵活、快速。该库提供了流行模型架构的 Keras 3 实现，并搭配了在 Kaggle 上提供的一系列预训练检查点。模型可以在 TensorFlow、Jax 和 Torch 任何一种后端上用于训练和推理。
KerasHub 是核心 Keras API 的扩展；KerasHub 组件以 keras.Layer 和 keras.Model 的形式提供。如果你熟悉 Keras，恭喜你！你已经理解了 KerasHub 的大部分内容。
本指南旨在成为该库的易懂介绍。我们将从使用高级 API 进行图像分类和文本生成开始，然后逐步展示模型的更深层定制和训练。在整个指南中，我们使用 Professor Keras（Keras 的官方吉祥物）作为材料复杂度的视觉参考。
一如既往，我们的 Keras 指南将专注于真实世界的代码示例。你可以随时通过点击指南顶部的 Colab 链接来运行这里的代码。
安装和设置
首先，我们来安装 keras-hub。该库在 PyPI 上可用，因此我们只需使用 pip 安装它即可。
!pip install --upgrade --quiet keras-hub keras
Keras 3 构建在 TensorFlow、Jax 和 Torch 后端之上。在编写 Keras 代码时，你应该在任何库导入之前首先指定后端。本指南我们将使用 Jax 后端，但你也可以使用 torch 或 tensorflow，而无需更改本指南其余部分的任何一行代码。这就是 Keras 3 的强大之处！
我们还将设置 XLA_PYTHON_CLIENT_MEM_FRACTION，这会从一开始就释放整个 GPU 供 Jax 使用。
import os
os.environ["KERAS_BACKEND"] = "jax" # or "tensorflow" or "torch"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.0"
最后，我们需要做一些额外的设置来访问本指南中使用的模型。许多流行的开放 LLM，例如 Google 的 Gemma 和 Meta 的 Llama，在访问模型权重之前需要接受社区许可。本指南中我们将使用 Gemma，因此可以按照以下步骤操作：
- 前往 Gemma 2 模型页面，并接受顶部横幅中的许可。
- 通过前往 Kaggle 设置并在“API”部分下点击“Create New Token”按钮来生成 Kaggle API 密钥。
- 在你的 Colab 笔记本内，点击左侧工具栏上的密钥图标。添加两个 secret：
KAGGLE_USERNAME（你的用户名）和KAGGLE_KEY（你刚刚创建的 API 密钥）。将这些 secret 设置为对你正在运行的笔记本可见。
API 快速入门
在我们开始之前，先看看 KerasHub 库中我们将要使用的关键类。
- 任务（Task）：例如，
keras_hub.models.CausalLM、keras_hub.models.ImageClassifier和keras_hub.models.TextClassifier。- 作用：任务将原始图像、音频和文本输入映射到模型预测。
- 重要性：任务是 KerasHub API 的最高级入口点。它将预处理和建模封装到一个简单易用的类中。任务可用于微调和推理。
- 包含：
backbone和preprocessor。 - 继承自：
keras.Model。
- 骨干网络（Backbone）：
keras_hub.models.Backbone。- 作用：将预处理后的张量输入映射到模型的潜在空间。
- 重要性：骨干网络以一种不特定于任何特定任务的方式封装了预训练模型的架构和参数。骨干网络可以与任意的预处理和将密集特征映射到预测的“头部”层结合，以完成任何机器学习任务。
- 继承自：
keras.Model。
- 预处理器（Preprocessor）：例如，
keras_hub.models.CausalLMPreprocessor、keras_hub.models.ImageClassifierPreprocessor和keras_hub.models.TextClassifierPreprocessor。- 作用：预处理器将原始图像、音频和文本输入映射到预处理后的张量输入。
- 重要性：预处理层封装了所有任务特定的预处理，例如图像大小调整和文本分词，可以独立使用来预计算预处理后的输入。请注意，如果你使用的是高级任务类，这种预处理默认已经内置了。
- 包含：
tokenizer、audio_converter和/或image_converter。 - 继承自：
keras.layers.Layer。
- 分词器（Tokenizer）：
keras_hub.tokenizers.Tokenizer。- 作用：将字符串转换为 token ID 序列。
- 重要性：字符串的原始字节是文本输入的低效表示形式，因此我们首先将字符串输入映射到整数 token ID。这个类封装了字符串到整数的映射以及反向映射（通过
detokenize()方法）。 - 继承自：
keras.layers.Layer。
- 图像转换器（ImageConverter）：
keras_hub.layers.ImageConverter。- 作用：调整和缩放图像输入。
- 重要性：图像模型通常需要将图像输入归一化到特定范围，或将输入调整到特定大小。这个类封装了特定于图像的预处理。
- 继承自：
keras.layers.Layer。
- 音频转换器（AudioConveter）：
keras_hub.layers.AudioConveter。- 作用：将原始音频转换为模型准备好的输入。
- 重要性：音频模型通常需要在将原始音频输入传递给模型之前对其进行预处理，例如计算音频信号的频谱图。这个类以易于使用的层封装了特定于图像的预处理。（译注：原文AudioConveter应为AudioConverter）
- 继承自：
keras.layers.Layer。
此处列出的所有类都有一个 from_preset() 构造函数，它将使用给定预训练模型标识符的权重和状态实例化组件。例如，keras_hub.tokenizers.Tokenizer.from_preset("gemma2_2b_en") 将创建一个使用 Gemma2 分词器词汇表对文本进行分词的层。
下图显示了所有这些核心类如何交互。箭头表示组合而非继承（例如，一个任务拥有一个骨干网络）。
图像分类
设置完成！让我们开始使用预训练模型。我们加载一张加州鹌鹑的测试图像并对其进行分类。
import keras
import numpy as np
import matplotlib.pyplot as plt
image_url = "https://upload.wikimedia.org/wikipedia/commons/a/aa/California_quail.jpg"
image_path = keras.utils.get_file(origin=image_url)
image = keras.utils.load_img(image_path)
plt.imshow(image)
我们可以使用在 ImageNet-1k 数据库上训练的 ResNet 视觉模型。该模型将为每个输入样本输出一个范围在 [0, 1000) 的标签，其中每个标签对应于一些现实世界的实体，如“奶罐”或“豪猪”。实际上，该数据集在索引 85 处有一个特定的鹌鹑标签。让我们下载模型并预测标签。
import keras_hub
image_classifier = keras_hub.models.ImageClassifier.from_preset(
"resnet_50_imagenet",
activation="softmax",
)
batch = np.array([image])
image_classifier.preprocessor.image_size = (224, 224)
preds = image_classifier.predict(batch)
preds.shape
1/1 ━━━━━━━━━━━━━━━━━━━━ 2s 2s/step
(1, 1000)
这些 ImageNet 标签并不是特别“人类可读”，因此我们可以使用内置的实用函数将预测解码为一组类名。
keras_hub.utils.decode_imagenet_predictions(preds)
[[('quail', 0.9996534585952759),
('prairie_chicken', 8.45497488626279e-05),
('partridge', 1.4000976079842076e-05),
('black_grouse', 7.407367775158491e-06),
('bullet_train', 7.323932550207246e-06)]]
看起来不错！模型权重成功下载，我们以几乎确定的概率预测了鹌鹑图像的正确分类标签。
这是我们第一次使用上面 API 快速入门中提到的高级任务（task） API 的例子。keras_hub.models.ImageClassifier 是一个用于图像分类的任务，可以与许多不同的模型架构（ResNet、VGG、MobileNet 等）一起使用。你可以在 Kaggle 上查看 Keras 团队直接提供的完整模型列表。
任务只是 keras.Model 的子类——你可以像任何其他模型一样在 classifier 对象上使用 fit()、compile() 和 save()。但 KerasHub 库为任务提供了一些额外功能。第一个也是最重要的是 from_preset()，这是你在 KerasHub 的许多类上都能看到的特殊构造函数。
预设（preset）是模型状态的目录。它定义了我们应该加载的架构以及与之匹配的预训练权重。from_preset() 允许我们从许多不同位置加载预设目录：
- 本地目录。
- Kaggle 模型中心。
- HuggingFace 模型中心。
你可以查看 keras_hub.models.ImageClassifier.from_preset 文档，以更好地理解从预设构建 Keras 模型时的所有选项。
所有任务都使用两个主要子对象。一个 keras_hub.models.Backbone 和一个 keras_hub.layers.Preprocessor。你可能已经熟悉计算机视觉中的骨干网络（backbone）一词，它通常用于描述将图像映射到潜在空间的特征提取网络。KerasHub 骨干网络是此概念的泛化，我们用它来指代任何没有任务特定头部的预训练模型。也就是说，KerasHub 骨干网络将原始图像、音频和文本（或这些输入的组合）映射到预训练模型的潜在空间。然后，我们可以将此潜在空间映射到任何数量的任务特定输出，具体取决于我们要使用模型做什么。
预处理器（preprocessor）只是一个 Keras 层，它执行特定任务的所有预处理。在我们的例子中，预处理将使用一些 ImageNet 特定的均值和方差数据来调整我们的输入图像大小并将其缩放到 [0, 1] 范围。让我们依次调用任务的预处理器和骨干网络，看看我们的输入形状会发生什么变化。
print("Raw input shape:", batch.shape)
resized_batch = image_classifier.preprocessor(batch)
print("Preprocessed input shape:", resized_batch.shape)
hidden_states = image_classifier.backbone(resized_batch)
print("Latent space shape:", hidden_states.shape)
Raw input shape: (1, 557, 707, 3)
Preprocessed input shape: (1, 224, 224, 3)
Latent space shape: (1, 7, 7, 2048)
我们的原始图像在预处理过程中被缩放到 (224, 224)，最后被缩小到一个 (7, 7) 图像的 2048 个特征向量——ResNet 模型的潜在空间。请注意，ResNet 实际上可以处理任意大小的图像，尽管如果你的图像大小与预训练数据差异很大，性能最终会下降。如果你想禁用预处理层中的大小调整，可以运行 image_classifier.preprocessor.image_size = None。
如果你想知道加载的任务的确切结构，可以像任何 Keras 模型一样使用 model.summary()。任务的模型摘要将包含有关模型预处理的额外信息。
image_classifier.summary()
Preprocessor: "res_net_image_classifier_preprocessor"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Config ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ res_net_image_converter │ Image size: (224, 224) │ │ (ResNetImageConverter) │ │ └──────────────────────────────────────────────┴───────────────────────────────┘
Model: "res_net_image_classifier"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input_layer (InputLayer) │ (None, None, None, 3) │ 0 │ ├───────────────────────────────────┼──────────────────────────┼───────────────┤ │ res_net_backbone (ResNetBackbone) │ (None, None, None, 2048) │ 23,561,152 │ ├───────────────────────────────────┼──────────────────────────┼───────────────┤ │ pooler (GlobalAveragePooling2D) │ (None, 2048) │ 0 │ ├───────────────────────────────────┼──────────────────────────┼───────────────┤ │ output_dropout (Dropout) │ (None, 2048) │ 0 │ ├───────────────────────────────────┼──────────────────────────┼───────────────┤ │ predictions (Dense) │ (None, 1000) │ 2,049,000 │ └───────────────────────────────────┴──────────────────────────┴───────────────┘
Total params: 25,610,152 (97.69 MB)
Trainable params: 25,557,032 (97.49 MB)
Non-trainable params: 53,120 (207.50 KB)
使用 LLM 生成文本
接下来，让我们尝试处理和生成文本。我们在生成文本时可以使用的任务是 keras_hub.models.CausalLM（其中 LM 是 Language Model 的缩写）。让我们下载 20 亿参数的 Gemma 2 模型并试用一下。
由于这个模型比我们刚刚下载的 ResNet 模型大了大约 100 倍，我们需要对 GPU 内存使用更加小心。我们可以使用半精度类型将我们约 25 亿参数的每个参数加载为两字节浮点数而不是四字节。为此，我们可以将 dtype 传递给 from_preset() 构造函数。from_preset() 会将任何 kwargs 转发给类的主构造函数，因此你可以传递适用于所有 Keras 层的 kwargs，例如 dtype、trainable 和 name。
causal_lm = keras_hub.models.CausalLM.from_preset(
"gemma2_instruct_2b_en",
dtype="bfloat16",
)
我们刚刚加载的模型是 Gemma 的指令微调版本，这意味着该模型经过了进一步的微调以用于聊天。只要我们坚持训练模型时使用的特定文本模板，就可以利用这些能力。这些特殊 token 因模型而异，可能很难追踪，Kaggle 模型页面将包含此类详细信息。
CausalLM 带有一个名为 generate() 的额外函数，可用于在循环中生成预测 token 并将其解码为字符串。
template = "<start_of_turn>user\n{question}<end_of_turn>\n<start_of_turn>model"
question = """Write a python program to generate the first 1000 prime numbers.
Just show the actual code."""
print(causal_lm.generate(template.format(question=question), max_length=512))
<start_of_turn>user
Write a python program to generate the first 1000 prime numbers.
Just show the actual code.<end_of_turn>
<start_of_turn>model
def is_prime(n):
if n <= 1:
return False
for i in range(2, int(n**0.5) + 1):
if n % i == 0:
return False
return True
count = 0
number = 2
primes = []
while count < 1000:
if is_prime(number):
primes.append(number)
count += 1
number += 1
print(primes)
<end_of_turn>
请注意，在 Jax 和 TensorFlow 后端，此 generate() 函数是编译的，因此当你第二次为相同的 max_length 调用它时，它实际上会快得多。KerasHub 将使用 Jax 和 TensorFlow 来计算一个可重用的优化版生成计算图。
question = "Share a very simple brownie recipe."
print(causal_lm.generate(template.format(question=question), max_length=512))
<start_of_turn>user
Share a very simple brownie recipe.<end_of_turn>
<start_of_turn>model
---
## Super Simple Brownies
**Ingredients:**
* 1 cup (2 sticks) unsalted butter, melted
* 2 cups granulated sugar
* 4 large eggs
* 1 teaspoon vanilla extract
* 1 cup all-purpose flour
* 1/2 cup unsweetened cocoa powder
* 1/4 teaspoon salt
**Instructions:**
1. Preheat oven to 350°F (175°C). Grease and flour a 9x13 inch baking pan.
2. In a large bowl, whisk together the melted butter and sugar until smooth.
3. Beat in the eggs one at a time, then stir in the vanilla extract.
4. In a separate bowl, whisk together the flour, cocoa powder, and salt.
5. Gradually add the dry ingredients to the wet ingredients, mixing until just combined. Do not overmix.
6. Pour the batter into the prepared pan and spread evenly.
7. Bake for 25-30 minutes, or until a toothpick inserted into the center comes out with a few moist crumbs attached.
8. Let cool completely before cutting and serving.
**Tips:**
* For extra fudgy brownies, underbake them slightly.
* Add chocolate chips, nuts, or other mix-ins to the batter for a personalized touch.
* Serve with a scoop of ice cream or whipped cream for a decadent treat.
Enjoy!
<end_of_turn>
与我们的图像分类器一样，我们可以使用模型摘要查看任务设置的详细信息，包括预处理。
causal_lm.summary()
Preprocessor: "gemma_causal_lm_preprocessor"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Config ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ gemma_tokenizer (GemmaTokenizer) │ Vocab size: 256,000 │ └──────────────────────────────────────────────┴───────────────────────────────┘
Model: "gemma_causal_lm"
┏━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━┩ │ padding_mask │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├───────────────────────┼───────────────────┼─────────────┼────────────────────┤ │ token_ids │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├───────────────────────┼───────────────────┼─────────────┼────────────────────┤ │ gemma_backbone │ (None, None, │ 2,614,341,… │ padding_mask[0][0… │ │ (GemmaBackbone) │ 2304) │ │ token_ids[0][0] │ ├───────────────────────┼───────────────────┼─────────────┼────────────────────┤ │ token_embedding │ (None, None, │ 589,824,000 │ gemma_backbone[0]… │ │ (ReversibleEmbedding) │ 256000) │ │ │ └───────────────────────┴───────────────────┴─────────────┴────────────────────┘
Total params: 2,614,341,888 (4.87 GB)
Trainable params: 2,614,341,888 (4.87 GB)
Non-trainable params: 0 (0.00 B)
我们的文本预处理包含一个分词器，这是所有 KerasHub 模型处理输入文本的方式。让我们直接使用它来更好地了解它是如何工作的。所有分词器都包含 tokenize() 和 detokenize() 方法，用于将字符串映射到整数序列以及将整数序列映射到字符串。直接调用层 tokenizer(inputs) 等同于调用 tokenizer.tokenize(inputs)。
tokenizer = causal_lm.preprocessor.tokenizer
tokens_ids = tokenizer.tokenize("The quick brown fox jumps over the lazy dog.")
print(tokens_ids)
string = tokenizer.detokenize(tokens_ids)
print(string)
[ 651 4320 8426 25341 36271 1163 573 27894 5929 235265]
The quick brown fox jumps over the lazy dog.
CausalLM 模型的 generate() 函数涉及采样步骤。Gemma 模型将为我们想要生成的每个 token 被调用一次，并返回所有 token 的概率分布。然后对该分布进行采样以选择序列中的下一个 token。
对于 Gemma 模型，我们默认使用贪婪采样，这意味着我们在每一步简单地选择模型中最可能的输出。但我们实际上可以使用所有 Keras 模型标准 compile 函数中的额外 sampler 参数来控制此过程。让我们来试试。
causal_lm.compile(
sampler=keras_hub.samplers.TopKSampler(k=10, temperature=2.0),
)
question = "Share a very simple brownie recipe."
print(causal_lm.generate(template.format(question=question), max_length=512))
<start_of_turn>user
Share a very simple brownie recipe.<end_of_turn>
<start_of_turn>model ## Ultimate Simple Brownies
This recipe requires NO oven or special equipment! Just microwave, mixing, and a few moments!
**Yields:** 6 large brownies
**Prep time:** 7 minutes
**Cook time:** 6-9 minutes, depending on your microwave
**What you need:**
* 3 ounces (about 2-3 tablespoons) chocolate chips
* 1/4 cup butter
* 1 large egg
* 1/2 cup granulated sugar
* 9 tablespoons all-purpose flour
**Optional Add-Ins (for extra fun):**
* 1/2 teaspoon vanilla
* 1/4 cup chopped walnuts or pecans
**Instructions:**
1. Place all microwave-safe mixing bowl ingredients:
- Chocolate Chips 🍫
- Butter 🧈
- Flour 🗲
- Egg (beaten!)
(You can add the optional add-INS like chopped nuts/extra vanilla, now is the good place to!)
2. Put all that in your microwave (microwave-safe dish or a heat-safe mug is fine!)
3. **Cook on:** Medium-high, stirring halfway.
* Time depends on your microwave, so keep checking, but aim for 6-9 minutes (if no stirring at least 8 mins). You want a thick, almost chewy-texture.
**To serve:** Cut up your brownies immediately and savor this classic treat. You'd also need a tall glass of cold milk or coffee (or both, if you've really enjoyed it).
Let me know if you want to experiment with a different chocolate or add-ins to make it even sweeter. Enjoy! 😉
<end_of_turn>
这里我们使用了 Top-K 采样器，这意味着我们将在每一步仅查看前 10 个预测 token 形成的偏分布中随机采样。我们还传递了一个 temperature 为 2，这在我们采样之前平坦化了我们的预测分布。
最终效果是，每次生成输出时，我们将更广泛地探索我们模型的分布。现在生成将是一个随机过程，每次我们重新运行生成时，都会得到不同的结果。我们可以注意到结果感觉比贪婪搜索更“随意”——有更多的小错误，一致性较低，以及将布朗尼放入微波炉的令人怀疑的建议。
你可以在 keras_hub.samplers 查看 Keras 支持的所有采样器。
在我们跳到下一节之前，让我们释放大型 Gemma 模型的内存。
del causal_lm
微调和发布图像分类器
现在我们已经尝试了图像和文本的推理，接下来尝试运行训练。我们将使用之前使用的 ResNet 图像分类器，并在简单的猫狗数据集上对其进行微调。我们可以先下载并提取数据。
import pathlib
extract_dir = keras.utils.get_file(
"cats_vs_dogs",
"https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_5340.zip",
extract=True,
)
data_dir = pathlib.Path(extract_dir) / "PetImages"
在处理大量真实世界图像数据时，损坏的图像是很常见的情况。让我们过滤掉其头部中不包含字符串“JFIF”的编码错误的图像。
num_skipped = 0
for path in data_dir.rglob("*.jpg"):
with open(path, "rb") as file:
is_jfif = b"JFIF" in file.peek(10)
if not is_jfif:
num_skipped += 1
os.remove(path)
print(f"Deleted {num_skipped} images.")
Deleted 1590 images.
我们可以使用 keras.utils.image_dataset_from_directory 加载数据集。这里需要注意的重要一点是，train_ds 和 val_ds 都将作为 tf.data.Dataset 对象返回，即使在 torch 和 jax 后端也是如此。
KerasHub 设计为使用 tf.data 作为在 CPU 上运行多线程预处理的默认 API。tf.data 是一个强大的 API，用于训练输入流水线，可以轻松扩展到复杂的、多主机训练作业。使用它不会限制你选择的后端，tf.data.Dataset 可以作为普通 numpy 数据的迭代器，并传递给任何 Keras 后端上的 fit()。
train_ds, val_ds = keras.utils.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="both",
seed=1337,
image_size=(256, 256),
batch_size=32,
)
Found 23410 files belonging to 2 classes.
Using 18728 files for training.
Using 4682 files for validation.
最简单地说，训练我们的分类器可以仅仅是调用模型的 fit() 方法并传入我们的数据集。但为了让这个例子更有趣，让我们展示如何在任务中定制预处理。
在第一个例子中，我们看到默认情况下，ResNet 模型的预处理会调整输入大小并重新缩放。这种预处理可以在我们创建模型时进行定制。我们可以使用 Keras 的图像预处理层来创建一个 keras.layers.Pipeline，它将对输入图像进行缩放、随机翻转和随机旋转。这些随机图像增强将使我们较小的数据集能够充当一个更大、更多样化的数据集。让我们来试试。
preprocessor = keras.layers.Pipeline(
[
keras.layers.Rescaling(1.0 / 255),
keras.layers.RandomFlip("horizontal"),
keras.layers.RandomRotation(0.2),
]
)
现在我们已经创建了一个新的预处理层，我们可以简单地在 from_preset() 构造函数期间将其传递给 ImageClassifier。我们还可以传递 num_classes=2 以匹配我们“猫”和“狗”的两个标签。当这样指定 num_classes 时，模型的头部权重将随机初始化，而不是包含 1000 个类别的图像分类权重。
image_classifier = keras_hub.models.ImageClassifier.from_preset(
"resnet_50_imagenet",
activation="softmax",
num_classes=2,
preprocessor=preprocessor,
)
请注意，如果你想在 Keras 之外预处理输入数据，只需将 preprocessor=None 传递给任务的 from_preset() 调用。在这种情况下，KerasHub 将完全不应用预处理，你可以使用任何库或工作流在将数据传递给 fit() 之前自由地预处理你的数据。
接下来，我们可以编译模型进行微调。KerasHub 任务只是一个普通的 keras.Model，带有一些额外功能，因此我们可以像分类任务一样正常进行 compile()。
image_classifier.compile(
optimizer=keras.optimizers.Adam(1e-4),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
至此，我们可以简单地运行 fit()。图像分类器将在训练模型时自动将我们的预处理应用于每个 batch。
image_classifier.fit(
train_ds,
validation_data=val_ds,
epochs=3,
)
Epoch 1/3
586/586 ━━━━━━━━━━━━━━━━━━━━ 0s 122ms/step - accuracy: 0.8869 - loss: 0.2921
Epoch 2/3
586/586 ━━━━━━━━━━━━━━━━━━━━ 65s 105ms/step - accuracy: 0.9858 - loss: 0.0393 - val_accuracy: 0.9912 - val_loss: 0.0234
Epoch 3/3
586/586 ━━━━━━━━━━━━━━━━━━━━ 57s 96ms/step - accuracy: 0.9897 - loss: 0.0289 - val_accuracy: 0.9930 - val_loss: 0.0206
<keras.src.callbacks.history.History at 0x787e77fb2550>
经过三个 epoch 的数据训练后，我们在猫狗验证数据集上取得了 99% 的准确率。这并不令人惊讶，考虑到我们开始使用的 ImageNet 预训练权重已经可以单独对一些品种的猫狗进行分类。
现在我们有一个微调后的模型，让我们尝试保存它。你只需运行 task.save_to_preset() 即可为任何任务创建并保存一个微调后的预设。
image_classifier.save_to_preset("cats_vs_dogs")
KerasHub 最强大的功能之一是将模型上传到 Kaggle 或 Huggingface 模型中心并与他人共享。keras_hub.upload_preset 允许你上传已保存的预设。
在这里，我们将上传到 Kaggle。我们已经通过 Kaggle 进行了认证，以便早些时候下载 Gemma 模型。运行以下单元格将把一个新模型上传到 Kaggle。
from google.colab import userdata
username = userdata.get("KAGGLE_USERNAME")
keras_hub.upload_preset(
f"kaggle://{username}/resnet/keras/cats_vs_dogs",
"cats_vs_dogs",
)
Uploading Model https://www.kaggle.com/models/matthewdwatson/resnet/keras/cats_vs_dogs ...
Upload successful: cats_vs_dogs/task.json (5KB)
Upload successful: cats_vs_dogs/task.weights.h5 (270MB)
Upload successful: cats_vs_dogs/metadata.json (157B)
Upload successful: cats_vs_dogs/model.weights.h5 (90MB)
Upload successful: cats_vs_dogs/config.json (841B)
Upload successful: cats_vs_dogs/preprocessor.json (3KB)
Your model instance version has been created.
Files are being processed...
See at: https://www.kaggle.com/models/matthewdwatson/resnet/keras/cats_vs_dogs
让我们看一下我们数据集中的一个测试图像。
image = keras.utils.load_img(data_dir / "Cat" / "6779.jpg")
plt.imshow(image)
如果我们等待几分钟，直到我们的模型上传在 Kaggle 端完成处理，我们就可以继续下载刚刚创建的模型，并用它来分类这个测试图像。
image_classifier = keras_hub.models.ImageClassifier.from_preset(
f"kaggle://{username}/resnet/keras/cats_vs_dogs",
)
print(image_classifier.predict(np.array([image])))
1/1 ━━━━━━━━━━━━━━━━━━━━ 2s 2s/step
[[9.999286e-01 7.135461e-05]]
恭喜你使用 KerasHub 上传了第一个模型！如果你想与他人分享你的工作，可以访问上传模型时打印出的模型链接，并在设置中将模型公开。
在进行本指南的最后一个示例之前，让我们删除此模型以释放内存。
del image_classifier
构建自定义文本分类器
作为本入门指南的最后一个示例，让我们看看如何使用更低级别的 Keras 和 KerasHub 组件构建自定义模型。我们将构建一个文本分类器，用于将 IMDb 数据集中的电影评论分类为正面或负面。
让我们下载数据集。
extract_dir = keras.utils.get_file(
"imdb_reviews",
origin="https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz",
extract=True,
)
data_dir = pathlib.Path(extract_dir) / "aclImdb"
IMDb 数据集包含大量未标注的电影评论。我们在这里不需要这些，可以直接删除它们。
import shutil
shutil.rmtree(data_dir / "train" / "unsup")
接下来，我们可以使用 keras.utils.text_dataset_from_directory 加载数据。与上面创建图像数据集类似，返回的数据集将是 tf.data.Dataset 对象。
raw_train_ds = keras.utils.text_dataset_from_directory(
data_dir / "train",
batch_size=2,
)
raw_val_ds = keras.utils.text_dataset_from_directory(
data_dir / "test",
batch_size=2,
)
Found 25000 files belonging to 2 classes.
Found 25000 files belonging to 2 classes.
KerasHub 被设计为一个分层 API。在最顶层，任务旨在使其易于快速解决问题。我们可以在这里继续使用任务 API，并为像 BERT 这样的文本分类模型创建一个 keras_hub.models.TextClassifer，然后在大约 10 行代码中对其进行微调。
相反，为了让我们的最后一个例子更有趣一些，让我们展示如何使用更低级别的 API 组件来做一些库中没有直接内置的事情。我们将使用我们之前使用的 Gemma 2 模型，它通常用于生成文本，并修改它以输出分类预测。
使用生成模型进行分类的常见方法是在生成语境中继续使用它，通过向其提示评论和一个问题（"Is this review positive or negative?"）。但如果你想要与标签关联的实际概率分数，构建一个实际的分类器会更有用。
我们可以不通过 CausalLM 任务加载 Gemma 2 模型，而是加载两个更低级别的组件：一个骨干网络（backbone）和一个分词器（tokenizer）。就像我们目前使用的任务类一样，keras_hub.models.Backbone 和 keras_hub.tokenizers.Tokenizer 都具有用于加载预训练模型的 from_preset() 构造函数。如果你正在运行这段代码，你会注意到我们第二次使用模型时不必等待下载，权重文件在第一次使用模型时已本地缓存。
tokenizer = keras_hub.tokenizers.Tokenizer.from_preset(
"gemma2_instruct_2b_en",
)
backbone = keras_hub.models.Backbone.from_preset(
"gemma2_instruct_2b_en",
)
我们在本指南的第二个示例中看到了分词器的作用。我们可以用它来将字符串输入映射到 token ID，这种方式与 Gemma 模型的预训练权重匹配。
骨干网络将把 token ID 序列映射到模型潜在空间中的嵌入 token 序列。我们可以利用这种丰富的表示来构建分类器。
让我们首先定义一个自定义预处理例程。keras_hub.layers 包含一组建模和预处理层，其中包括一些用于 token 预处理的层。我们可以使用 keras_hub.layers.StartEndPacker，它将在每条评论的开头添加一个特殊的起始 token，在结尾添加一个特殊的结束 token，最后将每条评论截断或填充到固定长度。
如果我们将此与我们的 tokenizer 结合，我们可以构建一个预处理函数，该函数将输出形状为 (batch_size, sequence_length) 的 token ID batch。我们还应该输出一个填充掩码，标记哪些 token 是填充 token，以便以后可以从 Transformer 的注意力计算中排除这些位置。KerasNLP 中的大多数 Transformer 骨干网络都接受一个 "padding_mask" 输入。
packer = keras_hub.layers.StartEndPacker(
start_value=tokenizer.start_token_id,
end_value=tokenizer.end_token_id,
pad_value=tokenizer.pad_token_id,
sequence_length=None,
)
def preprocess(x, y=None, sequence_length=256):
x = tokenizer(x)
x = packer(x, sequence_length=sequence_length)
x = {
"token_ids": x,
"padding_mask": x != tokenizer.pad_token_id,
}
return keras.utils.pack_x_y_sample_weight(x, y)
定义好预处理后，我们可以简单地使用 tf.data.Dataset.map 将我们的预处理应用于输入数据。
train_ds = raw_train_ds.map(preprocess, num_parallel_calls=16)
val_ds = raw_val_ds.map(preprocess, num_parallel_calls=16)
next(iter(train_ds))
({'token_ids': <tf.Tensor: shape=(2, 256), dtype=int32, numpy=
array([[ 2, 94300, 1185, ... 0]],
dtype=int32)>,
'padding_mask': <tf.Tensor: shape=(2, 256), dtype=bool, numpy=
array([[ True, True, True, ... False]])>},
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([1, 0], dtype=int32)>)
在一个 25 亿参数的模型上运行微调比我们之前训练的图像分类器昂贵得多，原因很简单，这个模型是 ResNet 大小的 100 倍！为了加快速度，我们将训练数据的大小减少到原始大小的十分之一。当然，与完整训练相比，这会牺牲一些性能，但可以确保本指南中的运行速度。
train_ds = train_ds.take(1000)
val_ds = val_ds.take(1000)
接下来，我们需要在骨干网络模型上附加一个分类头部。一般来说，文本 Transformer 骨干网络将输出形状为 (batch_size, sequence_length, hidden_dim) 的张量。使用此输入进行分类的主要需求是在序列维度上进行池化，以便每个输入示例都有一个单独的特征向量。
由于 Gemma 模型是一个生成模型，信息仅在序列中从左向右传递。唯一能够“看到”整个电影评论输入的 token 表示是每条评论中的最后一个 token。我们可以编写一个简单的池化层来完成此操作——我们只需获取每个输入序列的最后一个非填充位置。编写这样的层没有特殊的过程，我们可以正常使用 Keras 和 keras.ops。
from keras import ops
class LastTokenPooler(keras.layers.Layer):
def call(self, inputs, padding_mask):
end_positions = ops.sum(padding_mask, axis=1, keepdims=True) - 1
end_positions = ops.cast(end_positions, "int")[:, :, None]
outputs = ops.take_along_axis(inputs, end_positions, axis=1)
return ops.squeeze(outputs, axis=1)
有了这个池化层，我们就可以编写我们的 Gemma 分类器了。KerasHub 中的所有任务和骨干网络模型都是函数式模型，因此我们可以轻松地操作模型结构。我们将对输入调用骨干网络，添加新的池化层，最后添加一个小型前馈网络，中间使用 "relu" 激活。让我们来试试。
inputs = backbone.input
x = backbone(inputs)
x = LastTokenPooler(
name="pooler",
)(x, inputs["padding_mask"])
x = keras.layers.Dense(
2048,
activation="relu",
name="pooled_dense",
)(x)
x = keras.layers.Dropout(
0.1,
name="output_dropout",
)(x)
outputs = keras.layers.Dense(
2,
activation="softmax",
name="output_dense",
)(x)
text_classifier = keras.Model(inputs, outputs)
text_classifier.summary()
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ │ padding_mask │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ token_ids │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ gemma_backbone │ (None, None, │ 2,614,341… │ padding_mask[0][… │ │ (GemmaBackbone) │ 2304) │ │ token_ids[0][0] │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ pooler │ (None, 2304) │ 0 │ gemma_backbone[0… │ │ (LastTokenPooler) │ │ │ padding_mask[0][… │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ pooled_dense │ (None, 2048) │ 4,720,640 │ pooler[0][0] │ │ (Dense) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ output_dropout │ (None, 2048) │ 0 │ pooled_dense[0][… │ │ (Dropout) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ output_dense │ (None, 2) │ 4,098 │ output_dropout[0… │ │ (Dense) │ │ │ │ └─────────────────────┴───────────────────┴────────────┴───────────────────┘
Total params: 2,619,066,626 (9.76 GB)
Trainable params: 2,619,066,626 (9.76 GB)
Non-trainable params: 0 (0.00 B)
在训练之前，我们应该使用最后一个技巧，以便在免费的 Colab GPU 上运行这段代码。从我们的模型摘要中可以看出，模型占用了近 10 GB 的空间。优化器在训练期间需要创建每个参数的多个副本，使得模型在训练期间的总空间接近 30 或 40 GB。
这会导致许多 GPU OOM（内存溢出）。一个有用的技巧是启用骨干网络上的 LoRA。LoRA 是一种冻结整个模型，并仅训练大型权重矩阵的低参数分解的方法。你可以在这篇Keras 示例中阅读更多关于 LoRA 的信息。让我们尝试启用它并重新打印摘要。
backbone.enable_lora(4)
text_classifier.summary()
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ │ padding_mask │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ token_ids │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ gemma_backbone │ (None, None, │ 2,617,270… │ padding_mask[0][… │ │ (GemmaBackbone) │ 2304) │ │ token_ids[0][0] │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ pooler │ (None, 2304) │ 0 │ gemma_backbone[0… │ │ (LastTokenPooler) │ │ │ padding_mask[0][… │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ pooled_dense │ (None, 2048) │ 4,720,640 │ pooler[0][0] │ │ (Dense) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ output_dropout │ (None, 2048) │ 0 │ pooled_dense[0][… │ │ (Dropout) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ output_dense │ (None, 2) │ 4,098 │ output_dropout[0… │ │ (Dense) │ │ │ │ └─────────────────────┴───────────────────┴────────────┴───────────────────┘
Total params: 2,621,995,266 (9.77 GB)
Trainable params: 7,653,378 (29.20 MB)
Non-trainable params: 2,614,341,888 (9.74 GB)
启用 LoRA 后，我们的模型从 10GB 可训练参数减少到仅 20MB。这意味着优化器变量占用的空间不再是问题。
设置好这一切后，我们可以像往常一样编译和训练模型。
text_classifier.compile(
optimizer=keras.optimizers.Adam(5e-5),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
text_classifier.fit(
train_ds,
validation_data=val_ds,
)
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 295s 285ms/step - accuracy: 0.7733 - loss: 0.6511 - val_accuracy: 0.9370 - val_loss: 0.2814
<keras.src.callbacks.history.History at 0x787e103ae010>
我们在电影评论情感分类问题上取得了超过 ~93% 的准确率。这还不错，考虑到我们只使用了原始数据集的十分之一进行训练。
总而言之，我们在本例中创建的 backbone 和 tokenizer 使我们能够充分利用预训练的 Gemma 检查点，而不会限制我们可以用它们做什么。这是 KerasHub API 的核心目标。简单的工作流程应该容易实现，当你深入研究时，你将获得一套高度可定制的构建块。
进一步探索
这只是 KerasHub 功能的冰山一角。
本指南展示了 KerasHub 库附带的一些高级任务，但我们此处未涵盖的任务还有很多。例如，尝试使用 KerasHub 中的 Stable Diffusion 生成图像。
KerasHub 最显著的优势在于它为你提供了将预训练构建块与 Keras 3 的全部功能结合的灵活性。你可以使用 keras.distribution API 在 TPU 上通过模型并行化训练大型 LLM。你可以使用 Keras 的 量化方法对模型进行量化。你可以编写自定义训练循环，甚至可以混入直接的 Jax、Torch 或 Tensorflow 调用。
请访问 keras.io/keras_hub 查看完整的指南和示例列表，以继续深入研究该库。