保存Tensorflow_probability模型报cannot pickle module对象TypeError错误

这个报错是TensorFlow Probability（以下简称TFP）和Keras原生保存逻辑的已知兼容问题，不是代码编写错误。
你模型里用到的DenseFlipout是TFP提供的贝叶斯自定义层，不属于Keras原生内置层。当你直接调用model.save()尝试保存完整模型时，Keras默认会走pickle序列化逻辑处理自定义层对象，而TFP的贝叶斯层内部持有概率计算相关的模块引用，这类模块对象本身不支持pickle序列化，就会抛出TypeError: cannot pickle 'module' object的错误。
如果你的TFP版本在0.15以下，这个问题触发概率极高，属于长期存在的兼容缺陷，并非个别环境问题。

按稳定性从高到低排序：

• 仅保存模型权重，绕开整模型序列化逻辑
  放弃用model.save()存完整模型结构+权重+训练状态的方式，改用权重存储接口：

  # 保存权重
  best_model.save_weights("bnn_weights.h5")
  

  后续需要加载模型时，先手动复现和训练时完全一致的模型结构（层顺序、层参数都要完全匹配），再加载权重即可：

  # 先重新搭建和训练时完全相同的模型结构
  model = tf.keras.Sequential([
      tfp.layers.DenseFlipout(47, name="fcl_0"),
      tf.keras.layers.BatchNormalization(name="BN_0"),
      tfp.layers.DenseFlipout(14, name="fcl_1"),
      tfp.layers.DenseFlipout(1, name="final_layer")
  ])
  # 加载权重
  model.load_weights("bnn_weights.h5")
  

  这个方案完全绕开了pickle序列化自定义层的步骤，所有TF/TFP版本都能稳定使用，是最推荐的方案。
• 调整保存参数，用SavedModel格式跳过优化器存储
  如果你确实需要存完整模型，可以在调用save时强制指定用TF原生SavedModel格式，同时关闭优化器状态保存（优化器状态里也常带无法序列化的对象）：

  best_model.save(
      "model_BNN",
      save_format="tf",
      include_optimizer=False
  )
  

  加载时需要显式声明自定义层对象：

  model = tf.keras.models.load_model(
      "model_BNN",
      custom_objects={"DenseFlipout": tfp.layers.DenseFlipout}
  )
  

  注意这个方法不是所有版本都生效，加载后必须跑一次前向推理校验，确认贝叶斯层输出符合预期，避免出现参数丢失的问题。
• 升级TensorFlow和TFP到匹配的稳定版本
  0.15之后的TFP版本对Keras保存逻辑做了部分适配，如果你当前用的是比较老的版本，可以将TensorFlow和TFP升级到大版本号匹配的最新稳定版（比如TensorFlow 2.15匹配TFP 0.23），部分场景下不需要修改保存代码也能正常调用model.save()。

避坑提示：不要尝试用pickle直接dump整个模型，哪怕绕开了module对象的报错，存下来的模型也大概率会出现权重丢失、贝叶斯层随机种子失效的问题。

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