Функциональный API против создания подклассов объектно-ориентированного вариационного автоэнкодера

Question

Во-первых, я написал вариационный автоэнкодер через Functional API, используя тензорный поток v2.3, следуя этому примеру tf :

input_dim = X_train.shape[1:]
layer_dim = 128

# Build encoder
inputs = Input(shape = input_dim)
z1 = Flatten()(inputs)
z = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(z1)
z = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(z)
c1 = Concatenate()([z,z1])
z = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(c1)
z = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(z)
c2 = Concatenate()([z,c1])

# Latent layers and sampling
z_mean = Dense(z_size)(c2)
z_log_var = Dense(z_size)(c2)
z = Sampling(beta)([z_mean, z_log_var])

variational_encoder = Model(inputs=[inputs], outputs=[z_mean, z_log_var, z], name='Encoder')

# Build decoder
decoder_inputs = Input(shape=[z_size])
x = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(decoder_inputs)
x = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(x)
c1 = Concatenate()([x,decoder_inputs])
x = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(c1)
x = Dense(layer_dim, activation='elu', kernel_initializer=he_uniform())(x)
c2 = Concatenate()([x,c1])
x = Dense(np.multiply(*input_dim), activation='linear')(c2)
outputs = Reshape(input_dim)(x)

variational_decoder = Model(inputs=[decoder_inputs], outputs=[outputs], name='Decoder')

# Latent space and decoded latent space
_, _, codings = variational_encoder(inputs)
reconstructions = variational_decoder(codings)

# Build autoencoder
variational_ae = Model(inputs=[inputs], outputs=[reconstructions], name='VAE')

# Mean Squared Error loss
def loss(inputs, reconstructions):
    mse_loss = np.multiply(*input_dim) * mse(K.flatten(inputs), K.flatten(reconstructions))
    return (1.-beta) * K.mean(mse_loss)

def rounded_accuracy(y_true, y_pred):
    return binary_accuracy(tf.round(y_true), tf.round(y_pred))


variational_ae.compile(loss = loss,
                        optimizer = Adam())

history = variational_ae.fit(X_train, X_train, 
                            epochs = N_epochs,
                            batch_size = batch_size,
                            shuffle = True,
                            validation_split = 0.2,
                            callbacks=[checkpointer])

где выборка определяется как

class Sampling(Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""
    def __init__(self, beta):
        super(Sampling, self).__init__()
        self.beta = beta

    def call(self, inputs):
        mean, log_var = inputs
        batch = tf.shape(mean)[0]
        dim = tf.shape(mean)[1]

        # Add Kullback-Leibler divergence
        kl_loss = - 0.5 * K.sum(1 + log_var - K.square(mean) - K.exp(log_var), axis=-1)
        self.add_loss(self.beta * K.mean(kl_loss))

        # Add sampling
        epsilon = K.random_normal(shape=(batch, dim))
        return epsilon * K.exp(0.5 * log_var) + mean

Он отлично работает и очень хорошо реконструирует ввод. Однако я попытался узнать что-то новое и немного очистить свой код и попытался использовать метод подкласса, снова на той же странице ref . Я попытался воспроизвести приведенный выше код, но теперь результаты намного хуже, и я не знаю, откуда берутся различия, единственное различие, которое я вижу, - это слой Input (я пытался добавить его, но результатов нет. изменение). Новый код:

class Encoder(Model):
    """Maps CaloLayerEnergyDistributions to a triplet (z_mean, z_log_var, z)."""

    def __init__(self, beta, latent_dim=128, intermediate_dim=128, name="Encoder", **kwargs):
        super(Encoder, self).__init__(name=name, **kwargs)
        self.intermediate = [Dense(intermediate_dim, activation='elu', kernel_initializer=he_uniform()) for layer in range(4)]
        self.flatten = Flatten()
        self.concatenate = Concatenate()
        self.mean = Dense(latent_dim)
        self.log_var = Dense(latent_dim)
        self.sampling = Sampling(beta)

    def call(self, inputs):
        # Build encoder
        z1 = self.flatten(inputs)
        z = self.intermediate[0](z1)
        z = self.intermediate[1](z)
        c1 = self.concatenate([z, z1])
        z = self.intermediate[2](c1)
        z = self.intermediate[3](z)
        c2 = self.concatenate([z, c1])

        # Latent layers and sampling
        z_mean = self.mean(c2)
        z_log_var = self.log_var(c2)
        z = self.sampling((z_mean, z_log_var))
        return z_mean, z_log_var, z


class Decoder(Model):
    """Converts z, the encoded digit vector, back into energy distributions."""

    def __init__(self, original_dim, intermediate_dim=128, name="Decoder", **kwargs):
        super(Decoder, self).__init__(name=name, **kwargs)
        self.intermediate = [Dense(intermediate_dim, activation='elu', kernel_initializer=he_uniform()) for layer in range(4)]
        self.concatenate = Concatenate()
        self.outputs = Dense(np.multiply(*original_dim), activation='linear')
        self.reshape = Reshape(original_dim)

    def call(self, inputs):
        # Build decoder
        x = self.intermediate[0](inputs)
        x = self.intermediate[1](x)
        c1 = self.concatenate([x, inputs])
        x = self.intermediate[2](c1)
        x = self.intermediate[3](x)
        c2 = self.concatenate([x, c1])
        x = self.outputs(c2)
        return self.reshape(x)


class VariationalAutoEncoder(Model):
    """Combines the encoder and decoder into an end-to-end model for training."""

    def __init__(
        self,
        beta,
        original_dim,
        intermediate_dim=128,
        latent_dim=128,
        name="AutoEncoder",
        **kwargs
    ):
        super(VariationalAutoEncoder, self).__init__(name=name, **kwargs)
        self.original_dim = original_dim
        self.encoder = Encoder(beta)
        self.decoder = Decoder(original_dim)

    def call(self, inputs):
        _, _, codings = self.encoder(inputs)
        reconstructed = self.decoder(codings)

        return reconstructed

encoder = Encoder(beta)
decoder = Decoder(original_dim=input_dim)
bVAE = VariationalAutoEncoder(beta, input_dim)
bVAE.compile(loss = mse_loss, optimizer=Adam())

history = bVAE.fit(X_train, X_train,
                   epochs=N_epochs,
                   batch_size=batch_size,
                   shuffle=True,
                   validation_split=0.2,
                   callbacks=[checkpointer])

(выборка и потери остались прежними). Вопрос немного длинный, но я надеюсь, что кто-нибудь сможет мне помочь.