[FIXED] tensorflow.keras.utils.plot_model not working correctly

Issue

I've just used the tf.keras.utils.plot_model and the following diagram appeared:

The right part of the diagram shows the input and output shapes with a vertical split, but it should be an horizontal split to fit with the "intput:" and "output:" cells at its left.

code line :

# Import the required libraries.
from tensorflow.keras.layers import Input, Conv2D, Dense, Flatten
from tensorflow.keras.utils import plot_model
from tensorflow.keras.models import Model

input_shape = Input(shape=(128,128,3), name='input')
conv1 = Conv2D(filters=256, kernel_size=6, strides=2, padding='valid', activation='relu', name='conv1')(input_shape)
flt = Flatten()(conv1)
shared = Dense(64)(flt)
sub1 = Dense(16)(shared)
out1 = Dense(3, activation='softmax', name='gait')(sub1)
model = Model(inputs=input_shape, outputs=out1)


plot_model(model=model, show_shapes=True)

The result you want :

Solution

What you are looking for(Horizontal Split of Input and Output values) can be achieved in the earlier version of TensorFlow 2.7.

Please install the TensorFlow 2.7 using the below code, restart the kernel and then try running the same above code:

!pip install tensorflow==2.7

Answered By - TFer2