如何保存tf.keras模型各输入图像对应的整数预测类别

以下是问题中给出的基于MobileNetV2的5分类模型结构：

def locations_model(...):
    input_shape = image_shape + (3,)
    base_model = tf.keras.applications.MobileNetV2(...)
    base_model.trainable = False 
    inputs = tf.keras.Input(...)  
    # 中间特征提取、分类头结构省略
    outputs = tfl.Dense(5, activation = "softmax")(x)
    model = tf.keras.Model(inputs, outputs)
 
    return model

由于模型最后一层使用softmax激活函数，直接执行ypreds = model(inputs)或ypreds = model.predict(inputs)，只会得到一组和为1的5维数组，对应5个类别的预测概率，无法直接获取0-4区间整数形式的预测类别。

具体实现步骤 #

只需要在拿到预测概率后，沿类别维度取最大值对应的索引即可得到整数类别标签：

• 首先调用model.predict()遍历输入数据集，得到所有样本的预测概率矩阵
• 对概率矩阵的最后一维（即类别维度）调用argmax，返回的就是每个样本概率最高的类别对应的0-4整数标签

可直接参考以下代码：

# 得到测试集所有样本的预测概率，形状为(样本数, 5)
test_probas = model.predict(test_data)
# 沿最后一维取概率最大值的索引，得到形状为(样本数,)的整数类别数组
test_classes = test_probas.argmax(axis=-1)

得到test_classes数组后，可以直接用于后续精度计算、结果分析，也可以调用numpy的保存接口将结果持久化存储到本地。

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