TensorFlow加载模型权重后预测类别顺序错乱问题排查

模型实现 #

自定义Keras模型结构如下，基类封装预测阶段预处理逻辑，子类实现训练逻辑、自定义保存/加载逻辑：

class MyAbstractModel(keras.Model):
   ...
   def predict(self, **kwargs):
      kwargs["x"] = self.preprocess(kwargs["x"])
      super().predict(kwargs)

class MyModel(MyAbstractModel):
   def fit(self, **kwargs):
      kwargs["x"] = self.preprocess(kwargs["x"])
      kwargs["y"] = self.preprocess2(kwargs["y"])
      super().predict(kwargs)

   def save(self, path):
      # 为保留类内自定义逻辑，仅保存权重而非完整模型结构
      self.save_weights(path / "weights", save_format="tf")
      with open(path / "vectorizer", "wb") as file:
         joblib.dump({
            "config": self.vecorizer.get_config(),
            "weights": self.vecorizer.get_weights()
         }, file)  # self.vectorizer 为 keras.layers.TextVectorizer 类型层

   @classmethod
   def load(cls, path):
      instance = cls()
      instance.load_weights(path / "weights")
      with open(path / "vectorizer", "rb") as file:
         loaded_vec = joblib.load(file)
      self.vectorizer.from_config(loaded_vec["config"])
      self.vectorizer.set_weights(loaded_vec["weights"])
      return instance

• 模型为多分类任务设计，分类头为keras.layers.Dense(N_CLASSES, activation="softmax", name="classifier")
• 训练阶段通过keras.utils.to_categorical(label_mapping[y] for y in response)完成标签编码，其中label_mapping为类别到索引的映射字典，示例为{"A": 0, "B": 1, "C": 2}

故障现象 #

模型训练完成后，在同一会话内调用预测接口可得到正确结果，但保存后重新加载模型，预测返回的概率数组与类别的对应关系完全混乱：predict返回的N维（N为类别总数）概率数组，维度对应的类别顺序每次重载后都不一致。

• 曾尝试参考GitHub issue方案，在训练脚本中添加manual_variable_initialization(True)，问题未解决
• 加载模型时会输出如下警告，后续伴随一系列优化器状态未恢复的提示（如Value in checkpoint could not be found in the restored object: (root).optimizer's state 'm' for (root).classifier.kernel），经排查该类警告为TensorFlow已知问题，与本次类别顺序错乱故障无关：

WARNING:tensorflow: Detecting that an object or model or tf.train.Checkpoint is being deleted with unrestored values.

排查与根因 #

• 第一次尝试修复：将自定义预处理逻辑移至模型类外，分别使用joblib、keras.save_model保存完整模型对象，问题仍然存在。
• 最终定位：问题与Keras、TensorFlow框架无关，为业务代码逻辑错误导致。故障根源为类别列表初始化逻辑存在缺陷，导致预测结果映射类别时顺序不固定。
  原类别初始化代码如下：

class Category(enum.Enum):
   A = "A"
   B = "B"
   C = "C"

CATEGORIES = list(map(lambda c: c.value, Category))

通过上述方式从枚举类Category生成CATEGORIES类别列表时，不同运行进程中枚举值的遍历顺序不固定，最终导致概率数组维度与类别的映射关系每次重载都发生错乱。修复类别列表的固定顺序逻辑后，问题已解决。

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