Python中串联Keras Sequential模型异常：VAE结构汇总不符预期

问题描述 #

构建变分自动编码器（VAE）时，将两个Keras Sequential模型（Encoder和Decoder）串联后，VAE的结构汇总（summary）仅完整展示Encoder的层级，Decoder的所有层级被合并为一个Sequential块，未展开显示。

相关代码 #

import tensorflow as tf
from tensorflow import keras as tfk
from keras import layers as tfkl
import numpy as np
import keras_tuner as kt
import tensorflow_probability as tfp
from tensorflow_probability import layers as tfpl
from tensorflow_probability import distributions as tfd

def encoder_builder(hp, input_shape, latent_dim):

  prior = tfd.Independent(tfd.Normal(loc=tf.zeros(latent_dim), scale=1), reinterpreted_batch_ndims=1)

  enc = tfk.Sequential()
  enc.add(tfk.Input(shape=input_shape))
  enc.add(tfkl.Flatten())

  for i in range(hp.Int("enc_num_layers", 1, 3)):
    enc.add(
        tfkl.Dense(
            units=hp.Int(f"enc_units_{i}", min_value=32, max_value=512, step=32),
            activation=hp.Choice(f"enc_activation_{i}", ["relu", "tanh"]),
        )
    )

    enc.add(
        tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim), activation=None)
    )

    enc.add(
      tfpl.MultivariateNormalTriL(
          latent_dim,
          activity_regularizer=tfpl.KLDivergenceRegularizer(prior))
    )

    return enc

def decoder_builder(hp, input_shape, latent_dim):
    dec = tfk.Sequential()
    dec.add(tfk.Input(shape=latent_dim))
    for i in range(hp.Int("dec_num_layers", 2, 4)):
      dec.add(
          tfkl.Dense(
              units=hp.Int(f"dec_units_{i}", min_value=32, max_value=512, step=32),
              activation=hp.Choice(f"dec_activation_{i}", ["relu", "tanh"]),
          )
      )

    dec.add(tfkl.Dense(784, activation='relu'))
    dec.add(tfkl.Reshape((28, 28)))

    return dec

def model_builder(hp):
  encoder = encoder_builder(hp, (28, 28), 2)
  decoder = decoder_builder(hp, (28, 28), 2)

  encoder.summary()
  decoder.summary()

  vae = tfk.Model(inputs=encoder.inputs,
                outputs=decoder(encoder.outputs[0]))
  vae.summary()

结构汇总对比 #

编码器结构汇总 #

flatten_49 (Flatten)        (None, 784)               0
dense_190 (Dense)           (None, 32)                25120
dense_191 (Dense)           (None, 5)                 165
multivariate_normal_tri_l_45 (MultivariateNormalTriL)  ((None, 2), (None, 2))  0

解码器结构汇总 #

dense_192 (Dense)           (None, 32)                96
dense_193 (Dense)           (None, 32)                1056
dense_194 (Dense)           (None, 784)               25872
reshape_29 (Reshape)        (None, 28, 28)            0

VAE结构汇总（异常） #

input_65 (InputLayer)       [(None, 28, 28)]          0
flatten_49 (Flatten)        (None, 784)               0
dense_190 (Dense)           (None, 32)                25120
dense_191 (Dense)           (None, 5)                 165
multivariate_normal_tri_l_45 (MultivariateNormalTriL)  ((None, 2), (None, 2))  0
sequential_74 (Sequential)  (None, 28, 28)            27024

问题原因 #

Keras默认将嵌套的Sequential模型作为整体模块展示，不会自动展开内部层级。此外原代码中encoder_builder存在逻辑错误：return enc被放在了for循环内部，导致Encoder仅构建第一层就提前返回。

解决方案 #

1. 修正Encoder的逻辑错误 #

将return enc移到for循环外部，确保完整构建所有层级：

def encoder_builder(hp, input_shape, latent_dim):
    prior = tfd.Independent(tfd.Normal(loc=tf.zeros(latent_dim), scale=1), reinterpreted_batch_ndims=1)
    enc = tfk.Sequential()
    enc.add(tfk.Input(shape=input_shape))
    enc.add(tfkl.Flatten())

    for i in range(hp.Int("enc_num_layers", 1, 3)):
        enc.add(
            tfkl.Dense(
                units=hp.Int(f"enc_units_{i}", min_value=32, max_value=512, step=32),
                activation=hp.Choice(f"enc_activation_{i}", ["relu", "tanh"]),
            )
        )
        enc.add(
            tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim), activation=None)
        )
        enc.add(
            tfpl.MultivariateNormalTriL(
                latent_dim,
                activity_regularizer=tfpl.KLDivergenceRegularizer(prior))
        )
    # 移到循环外部
    return enc

2. 展开Decoder层级到VAE中 #

有两种方式实现Decoder层级的展开显示：

方式一：直接在VAE中构建Decoder层级

def model_builder(hp):
    encoder = encoder_builder(hp, (28, 28), 2)
    z = encoder.outputs[0]
    
    # 直接构建Decoder层级
    for i in range(hp.Int("dec_num_layers", 2, 4)):
        z = tfkl.Dense(
            units=hp.Int(f"dec_units_{i}", min_value=32, max_value=512, step=32),
            activation=hp.Choice(f"dec_activation_{i}", ["relu", "tanh"]),
        )(z)
    z = tfkl.Dense(784, activation='relu')(z)
    output = tfkl.Reshape((28, 28))(z)
    
    vae = tfk.Model(inputs=encoder.inputs, outputs=output)
    vae.summary()

方式二：手动遍历Decoder的层并添加到VAE

def model_builder(hp):
    encoder = encoder_builder(hp, (28, 28), 2)
    decoder = decoder_builder(hp, (28, 28), 2)
    
    # 依次应用Decoder的每一层
    x = encoder.outputs[0]
    for layer in decoder.layers:
        x = layer(x)
    
    vae = tfk.Model(inputs=encoder.inputs, outputs=x)
    vae.summary()

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