How do I plot a Keras/Tensorflow subclassing API model?

by

I’ve found some workaround to plot with the model sub-classing API. For the obvious reason Sub-Classing API doesn’t support Sequential or Functional API like model.summary() and nice visualization using plot_model. Here, I will demonstrate both.

class my_model(keras.Model):
    def __init__(self, dim):
        super(my_model, self).__init__()
        self.Base  = keras.keras.applications.VGG16(
            input_shape=(dim), 
            include_top = False, 
            weights="imagenet"
        )
        self.GAP   = L.GlobalAveragePooling2D()
        self.BAT   = L.BatchNormalization()
        self.DROP  = L.Dropout(rate=0.1)
        self.DENS  = L.Dense(256, activation='relu', name="dense_A")
        self.OUT   = L.Dense(1, activation='sigmoid')
    
    def call(self, inputs):
        x  = self.Base(inputs)
        g  = self.GAP(x)
        b  = self.BAT(g)
        d  = self.DROP(b)
        d  = self.DENS(d)
        return self.OUT(d)
    
    # AFAIK: The most convenient method to print model.summary() 
    # similar to the sequential or functional API like.
    def build_graph(self):
        x = Input(shape=(dim))
        return Model(inputs=[x], outputs=self.call(x))

dim = (124,124,3)
model = my_model((dim))
model.build((None, *dim))
model.build_graph().summary()

It will produce as follows:

Layer (type)                 Output Shape              Param #   
=================================================================
input_67 (InputLayer)        [(None, 124, 124, 3)]     0         
_________________________________________________________________
vgg16 (Functional)           (None, 3, 3, 512)         14714688  
_________________________________________________________________
global_average_pooling2d_32  (None, 512)               0         
_________________________________________________________________
batch_normalization_7 (Batch (None, 512)               2048      
_________________________________________________________________
dropout_5 (Dropout)          (None, 512)               0         
_________________________________________________________________
dense_A (Dense)              (None, 256)               402192    
_________________________________________________________________
dense_7 (Dense)              (None, 1)                 785       
=================================================================
Total params: 14,848,321
Trainable params: 14,847,297
Non-trainable params: 1,024

Now by using the build_graph function, we can simply plot the whole architecture.

# Just showing all possible argument for newcomer.  
tf.keras.utils.plot_model(
    model.build_graph(),                      # here is the trick (for now)
    to_file="model.png", dpi=96,              # saving  
    show_shapes=True, show_layer_names=True,  # show shapes and layer name
    expand_nested=False                       # will show nested block
)

It will produce as follows: -)

a

Similar QnA:

  1. Retrieving Keras Layer Properties from a tf.keras.Model
  2. Visualize nested keras.Model (SubClassed API) GAN model

More Related Contents:

  1. Keras, How to get the output of each layer?
  2. How to apply gradient clipping in TensorFlow?
  3. Keras replacing input layer
  4. Keras custom loss function: Accessing current input pattern
  5. How can I use a pre-trained neural network with grayscale images?
  6. Dimension of shape in conv1D
  7. Removing then Inserting a New Middle Layer in a Keras Model
  8. Custom loss function in Keras
  9. How do I get the weights of a layer in Keras?
  10. Multiple outputs in Keras
  11. ValueError: Input 0 of layer sequential is incompatible with the layer: : expected min_ndim=4, found ndim=3. Full shape received: [8, 28, 28]
  12. Keras LSTM input dimension setting
  13. How to implement Grad-CAM on a trained network
  14. How to fix “ResourceExhaustedError: OOM when allocating tensor”
  15. Strange behaviour of the loss function in keras model, with pretrained convolutional base
  16. Multi-input Multi-output Model with Keras Functional API
  17. How to make a custom activation function with only Python in Tensorflow?
  18. Deep-Learning Nan loss reasons
  19. Failed to get convolution algorithm. This is probably because cuDNN failed to initialize,
  20. Why is TensorFlow 2 much slower than TensorFlow 1?
  21. Can Keras with Tensorflow backend be forced to use CPU or GPU at will?
  22. What is the difference between tf.keras and tf.python.keras?
  23. Negative dimension size caused by subtracting 3 from 1 for ‘conv2d_2/convolution’
  24. How to know if underfitting or overfitting is occuring?
  25. How to handle non-determinism when training on a GPU?
  26. ‘Tensor’ object has no attribute ‘lower’
  27. ValueError: Tensor must be from the same graph as Tensor with Bidirectinal RNN in Tensorflow
  28. ValueError: Output tensors to a Model must be the output of a TensorFlow `Layer`
  29. What is right batch normalization function in Tensorflow?
  30. TensorFlow 2.0 Keras: How to write image summaries for TensorBoard

Leave a Comment