如何克隆/复制含TensorFlow Probability层的神经网络模型？

环境信息 #

• Python 3.8.11
• TensorFlow 2.10.0
• TensorFlow Probability 0.18.0

问题背景 #

模型构建使用了TFP的DenseVariational层，结构如下：

def posterior_mean_field(kernel_size: int, bias_size: int, dtype: Any) -> tf.keras.Model:
    n = kernel_size + bias_size
    c = np.log(np.expm1(1.))

    return tf.keras.Sequential([
        tfp.layers.VariableLayer(2 * n, dtype=dtype),
        tfp.layers.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t[..., :n],
                                                                           scale=1e-5 + tf.nn.softplus(c + t[..., n:])),
                                                                reinterpreted_batch_ndims=1)),
    ])


def prior_trainable(kernel_size: int, bias_size: int, dtype: Any) -> tf.keras.Model:
    n = kernel_size + bias_size
    return tf.keras.Sequential([
        tfp.layers.VariableLayer(n, dtype=dtype),
        tfp.layers.DistributionLambda(lambda t: tfd.Independent(
            tfd.Normal(loc=t, scale=1),
            reinterpreted_batch_ndims=1)),
    ])


def build_model():
    model = keras.Sequential([
        tfp.layers.DenseVariational(64, activation='relu', input_shape=[len(train_dataset.keys())],
                                    make_posterior_fn=posterior_mean_field,
                                    make_prior_fn=prior_trainable),
        layers.Dense(64, activation='relu'),
        layers.Dense(1),
    ])
    optimizer = tf.keras.optimizers.RMSprop(0.001)
    model.compile(loss='mse', optimizer=optimizer, metrics=['mae', 'mse'])
    return model


model = build_model()
model.build((3, 10))

尝试使用clone_model克隆模型时，触发错误：

NotImplementedError: Layer DenseVariational has arguments in `__init__` and therefore must override `get_config`.

移除DenseVariational层后，clone_model可正常工作，但保留该层时无法使用此方法。

解决方法 #

方法1：手动重建模型+复制权重 #

跳过clone_model，直接调用模型构建函数生成新模型，再复制权重。此方法不依赖层的序列化功能，只要新模型的层结构、输入形状与原模型完全一致即可：

# 重建模型结构
model_copy = build_model()
# 匹配原模型的输入形状
model_copy.build((3, 10))
# 复制原模型权重
model_copy.set_weights(model.get_weights())

方法2：自定义DenseVariational层并重写get_config #

如果一定要使用clone_model，可以继承原DenseVariational层，手动实现get_config方法。注意：由于make_posterior_fn和make_prior_fn是函数对象，序列化时可能存在限制，此方法仅适用于简单场景：

class CustomDenseVariational(tfp.layers.DenseVariational):
    def get_config(self):
        # 获取父类基础配置
        config = super().get_config()
        # 添加__init__中的自定义参数
        config.update({
            'make_posterior_fn': self.make_posterior_fn,
            'make_prior_fn': self.make_prior_fn,
        })
        return config

# 修改build_model函数，替换为自定义层
def build_model():
    model = keras.Sequential([
        CustomDenseVariational(64, activation='relu', input_shape=[len(train_dataset.keys())],
                               make_posterior_fn=posterior_mean_field,
                               make_prior_fn=prior_trainable),
        layers.Dense(64, activation='relu'),
        layers.Dense(1),
    ])
    optimizer = tf.keras.optimizers.RMSprop(0.001)
    model.compile(loss='mse', optimizer=optimizer, metrics=['mae', 'mse'])
    return model

是否需要提交GitHub Issue #

是的，建议在TensorFlow Probability的GitHub仓库提交Issue。DenseVariational层未实现get_config方法，导致无法兼容Keras的clone_model等依赖层序列化的功能，这属于官方层的功能缺失，提交Issue有助于官方修复该问题。

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