如何让Keras多分类模型调用predict时直接输出类别名而非概率？

核心思路是训练阶段保持原有流程不变，推理阶段将标签映射、概率取值的后处理逻辑嵌入模型内部，封装为端到端的推理模型，不需要修改训练逻辑，也不需要在模型外部维护编码器。

完整实现代码 #

1. 训练基础分类模型 #

训练阶段和你原有流程完全一致，只需要提前固定好所有类别的字符串列表即可：

import tensorflow as tf
import numpy as np

# 提前定义你的所有类别字符串，不需要依赖LabelEncoder
class_names = ["类别A", "类别B", "类别C"]
num_classes = len(class_names)

# 构建基础分类模型，结构和你原有模型完全一致，最后一层输出softmax概率
base_model = tf.keras.Sequential([
    # 这里替换为你的原有网络结构，示例为全连接网络
    tf.keras.layers.Dense(128, activation="relu", input_shape=(输入特征维度,)),
    tf.keras.layers.Dense(num_classes, activation="softmax")
])

# 训练逻辑和原有流程完全一致，还是用one-hot编码的标签训练
base_model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
# 省略base_model.fit(...)训练代码，训练完成后得到基础模型

2. 封装端到端推理模型 #

将后处理逻辑嵌入模型内部，用TensorFlow内置的哈希表实现索引到字符串的映射，支持模型序列化保存：

# 构建索引转类别字符串的静态哈希表，会随模型一起序列化
keys = tf.range(num_classes, dtype=tf.int64)
values = tf.constant(class_names, dtype=tf.string)
class_map_table = tf.lookup.StaticHashTable(
    initializer=tf.lookup.KeyValueTensorInitializer(keys, values),
    default_value="未知类别"
)

# 定义后处理逻辑：取最高概率+转类别字符串
def post_process(prob_matrix):
    max_prob = tf.reduce_max(prob_matrix, axis=-1)
    class_index = tf.argmax(prob_matrix, axis=-1)
    class_string = class_map_table.lookup(class_index)
    return class_string, max_prob

# 组装为完整的端到端推理模型
inputs = tf.keras.Input(shape=(输入特征维度,))
prob_output = base_model(inputs)
final_output = tf.keras.layers.Lambda(post_process)(prob_output)
infer_model = tf.keras.Model(inputs=inputs, outputs=final_output)

3. 推理调用效果 #

直接调用predict方法即可得到类别字符串和对应概率，不需要任何外部解码：

# 测试输入，替换为你的实际推理数据
test_input = np.random.random((1, 输入特征维度))
pred_class, pred_prob = infer_model.predict(test_input, verbose=0)

# 直接输出结果
print(pred_class[0].decode("utf-8")) # 输出类别字符串，如"类别B"
print(pred_prob[0]) # 输出对应概率值，如0.87

注意事项 #

• 封装好的infer_model可以直接用infer_model.save("xxx.h5")或者SavedModel格式保存，加载后可直接推理，不需要额外维护类别映射文件或编码器。
• 如果需要返回TopK个结果，只需修改post_process函数，用tf.math.top_k替换tf.argmax和tf.reduce_max即可。
• 训练阶段只用base_model即可，不需要对infer_model做训练操作。

内容的提问来源于stack exchange，提问作者ilI47oPurl1a