TensorFlow Probability自定义层模型保存后加载报错问询

问题原因 #

• 你只在custom_objects里注册了Conv3DReparameterization层，但模型中还使用了tfp.layers.DenseFlipout自定义层，未注册会导致加载时识别异常。
• 用HDF5格式（.h5后缀）保存包含TensorFlow Probability自定义层的模型时，常会出现参数序列化错误：比如将可调用的激活函数tf.nn.relu序列化为字符串，加载时尝试调用字符串对象就会抛出'str' object is not callable报错。

修复方案 #

方案1：补全custom_objects映射（适配.h5格式模型） #

加载时把所有用到的自定义对象都注册到custom_objects中即可：

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.models import load_model

custom_objects= {
    'Conv3DReparameterization': tfp.layers.Convolution3DReparameterization,
    'DenseFlipout': tfp.layers.DenseFlipout,
    'relu': tf.nn.relu
}
model3 = load_model('tf_test_model3.h5', custom_objects=custom_objects)

方案2：改用SavedModel格式保存加载（更稳定，推荐） #

SavedModel是TensorFlow官方推荐的序列化格式，对自定义层的兼容性远高于HDF5格式，不需要手动维护custom_objects映射：

保存模型时代码调整

# 直接传文件夹名，不要加.h5后缀
model3.save('tf_test_model3')

加载模型代码

model3 = tf.keras.models.load_model('tf_test_model3')

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