TF 2.6保存自编码器时AssignVariableOp类型不匹配错误如何解决

根因说明 #

你遇到的Input 'resource' of 'AssignVariableOp' Op has type float32 that does not match expected type of resource.错误，是因为TensorFlow 2.6版本强化了资源变量的类型校验逻辑，禁止了旧版本中对AssignVariableOp操作的隐式类型转换，通常出现在类别特征对应的嵌入层类型不匹配、自定义层变量序列化异常、使用HDF5格式保存自定义模型的场景。

修复步骤 #

• 改用TensorFlow官方推荐的SavedModel格式保存模型，弃用HDF5（.h5）格式。2.6版本后HDF5对自定义层、资源变量的兼容性大幅降低，保存和加载代码调整为：

  # 保存 不要加.h5后缀
  model.save("categorical_ae_model")
  # 加载
  loaded_model = tf.keras.models.load_model("categorical_ae_model")
  

• 修正类别特征的数据类型：输入嵌入层的类别编码数据必须为整数类型，不要转为float32：

  # 错误写法 给嵌入层传入浮点型输入
  x_train = x_train.astype("float32")
  # 正确写法 类别编码后统一转为int32
  x_train = x_train.astype("int32")
  x_test = x_test.astype("int32")
  

• 补全自定义层的序列化方法：如果自编码器用到自定义层，补充get_config和from_config方法保证变量能被正确识别为资源类型，而非普通张量：

  class CustomAELayer(tf.keras.layers.Layer):
      # 你的原有初始化、调用逻辑
      def get_config(self):
          config = super().get_config()
          # 追加层的自定义参数，示例如下，替换为你的实际参数
          config.update({
              "hidden_dim": self.hidden_dim,
              "dropout_rate": self.dropout_rate
          })
          return config
      @classmethod
      def from_config(cls, config):
          return cls(**config)
  

• 可选兼容配置：如果仍有报错，保存时关闭不必要的序列化选项：

  model.save(
      "categorical_ae_model",
      save_traces=False,
      include_optimizer=False # 不需要继续训练时添加
  )
  

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