Keras 'plot_model' shows wrong graph for nested models (Autoencoder)

859 Views Asked by At

When I create an autoencoder architecture with multiple inputs and outputs, the plot_model graph does not show up as expected (problems highlighted in red).

I assume the first issue occurs because I use encoder.inputs for the autoencoder. However, creating new input layers for the autoencoder results in an error for me (graph disconnected).

The problem is likely on my side and not a bug in Keras, so hopefully someone with more experience can guide me in the right direction.

ps. I can't just use a single autoencoder model, which would avoid this problem, because the same encoder and decoder models are additionally used in a GAN setting. (implementation is currently working with single inputs and slicing, but I'd really like to switch to multi inputs cause it feels way cleaner)

Code and architecture image below: enter image description here

from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
from tensorflow.keras.utils import plot_model

# Config encoder
state_inputs = Input(shape=7, name="encoder_state_inputs")
action_inputs = Input(shape=4, name="encoder_action_inputs")
encoder_layer = Dense(11, activation=LeakyReLU(alpha=0.2), name=f"encoder_layer")
classification_layer = Dense(4, activation="softmax", name=f"classification")
latent_layer = Dense(5, activation=LeakyReLU(alpha=0.2), name="latent_space")

# Build encoder
x = concatenate([state_inputs, action_inputs])
x = encoder_layer(x)
classifier = classification_layer(x)
latent = latent_layer(x)
encoder = Model([state_inputs, action_inputs], [classifier, latent], name="encoder")
plot_model(encoder, "encoder.png")


# Config decoder
classifier_inputs = Input(shape=4, name="classifier_inputs")
latent_inputs = Input(shape=5, name="latent_inputs")
decoder_layer = Dense(11, activation=LeakyReLU(alpha=0.2), name=f"decoder_layer")
state_reconstruction_layer = Dense(7, activation="softmax", name=f"state_reconstruction")
action_reconstruction_layer = Dense(4, activation="softmax", name=f"action_reconstruction")

# Build decoder
x = concatenate([classifier_inputs, latent_inputs])
x = decoder_layer(x)
state_reconstruction = state_reconstruction_layer(x)
action_reconstruction = action_reconstruction_layer(x)
decoder = Model([classifier_inputs, latent_inputs], [state_reconstruction, action_reconstruction], name="decoder")
plot_model(decoder, "decoder.png")

# Build autoencoder
encoded = encoder(encoder.inputs)
decoded = decoder(encoded)
autoencoder = Model(encoder.inputs, outputs=[classifier, decoded], name="autoencoder")

plot_model(autoencoder, "autoencoder.png", show_shapes=True, expand_nested=True)

1

There are 1 best solutions below

0
On

I was able to prevent the first problem by adding two input layers to the autoencoder.

The second problem (multiple outputs just connecting to one input) seems to be a known bug: https://github.com/tensorflow/tensorflow/issues/42101