Keras RaggedTensor 输入模型保存后怎么保持变长序列结构

我第一次把不等长的文本 token 送进 Keras 时，真正容易出错的地方不是模型层，而是输入签名：模型保存之后，如果加载出来的输入边界已经被当成普通 Tensor，原来的每条序列长度信息就无法按预期保留。这个小项目的做法是把 ragged=True 写进 Keras 输入，并明确使用 TensorFlow backend，再用原生 .keras 文件保存整个模型。

官方地址：https://keras.io/api/layers/core_layers/input/

先把动态长度和 RaggedTensor 分开

“序列长度不固定”和“输入是 RaggedTensor”是两个相关但不相同的概念。shape=(None,) 只说明这一维没有固定数值；如果每个 batch 样本长度都不同，还需要让输入对象知道这种分段结构。RaggedTensor 可以把不同长度的行放在同一个批次里，并通过 values 与 row_splits 保留边界。

Keras 当前的 Input API 将 ragged 标为 TensorFlow backend 支持的选项。我的取舍是先把 backend 写死在项目说明里，避免读者在 JAX 或 PyTorch backend 下照抄后才发现输入类型不兼容。

用一个最小模型接住变长序列

下面的模型只做一件事：把每条不同长度的整数序列映射成一个分类结果。Embedding 负责把 token 变成向量，GlobalAveragePooling1D 把变长序列汇总成固定宽度，最后由全连接层输出结果。这样可以把“变长”限制在输入和池化之前，保存问题更容易观察。

import os
import numpy as np
import tensorflow as tf
import keras

# RaggedTensor 输入目前绑定 TensorFlow backend；先显式声明，避免环境误切换。
os.environ["KERAS_BACKEND"] = "tensorflow"

# None 表示每条序列长度可变，ragged=True 保留每行的边界信息。
inputs = keras.Input(
    shape=(None,), dtype="int32", ragged=True, name="token_ids"
)
x = keras.layers.Embedding(input_dim=1000, output_dim=16, name="token_embedding")(inputs)
x = keras.layers.GlobalAveragePooling1D(name="sequence_pooling")(x)
outputs = keras.layers.Dense(2, activation="softmax", name="classifier")(x)
model = keras.Model(inputs, outputs, name="ragged_classifier")

# 使用不同长度的行组成一个 RaggedTensor，而不是硬塞进二维数组。
sample = tf.ragged.constant([[4, 8, 15], [16, 23], [42]], dtype=tf.int32)
prediction = model(sample, training=False)
print(prediction.shape)

这里的关键不是 None 本身，而是 ragged=True 和调用时的 tf.ragged.constant 必须配套。若把样本提前补齐成普通矩阵，模型仍可能运行，但那已经是 padding 方案，不再是本文讨论的输入结构。

用 .keras 保存模型，再按原结构加载

我更愿意保存完整的 Keras 模型，而不是只保存权重。这样输入签名、层配置和权重可以一起进入一个 .keras 文件，恢复时直接使用 keras.models.load_model。如果项目里还有自定义层，仍要按 Keras 的序列化规则注册或提供 custom_objects；RaggedTensor 并不会替你解决自定义对象的加载问题。

import keras
import tensorflow as tf

# 保存完整模型，扩展名使用 Keras 原生格式。
model.save("ragged_classifier.keras")

# 重新加载结构与权重；自定义层项目应在这里提供 custom_objects。
restored = keras.models.load_model("ragged_classifier.keras")

# 恢复后的模型继续接收每行长度不同的输入。
restored_sample = tf.ragged.constant([[7, 11], [13, 17, 19, 23]], dtype=tf.int32)
restored_prediction = restored(restored_sample, training=False)
print(restored_prediction.shape)

保存和加载负责恢复模型本身，输入数据仍由调用方构造。换句话说，.keras 文件不会把下一次请求的 RaggedTensor 数据一起保存进去；服务端接收请求后，依然要把 token 列表组装成正确的 ragged 结构。

这几个边界决定了恢复是否稳定

还有一个容易被忽略的选择：如果目标是跨 JAX、TensorFlow、PyTorch 复用同一套模型代码，RaggedTensor 输入就不是无条件成立的抽象。对我来说，项目如果确实依赖不等长批次，会把 TensorFlow backend 作为部署契约；如果更看重多后端迁移，则会在输入层明确选择 padding 或其他统一表示。

常见问题

shape=(None,) 能代替 ragged=True 吗？

不能完全代替。它只描述某个维度的动态形状；是否按每行不同长度保存边界，还取决于实际输入类型和模型输入的 ragged 设置。

重新加载后还需要重新声明 RaggedTensor 吗？

模型的输入签名会随完整模型保存，但每次推理的数据仍要按 RaggedTensor 构造。保存模型不等于保存未来请求的数据结构。

RaggedTensor 输入能在所有 Keras backend 使用吗？

不能把它当成跨后端通用能力。当前 Keras Input 文档明确把 ragged 支持限定在 TensorFlow backend；多后端项目应提前做输入表示取舍。

只保存权重可以吗？

可以，但加载前必须用完全匹配的代码重新创建模型。若希望输入签名、结构与权重一起恢复，直接保存完整的 .keras 模型更省心。

这次小项目里，稳定性的核心不是某个特殊保存参数，而是把“TensorFlow backend + ragged=True + RaggedTensor 调用数据”视为同一个契约。只要三者没有被拆开，模型保存和重新加载就不会把变长序列悄悄改成普通定长输入。