Keras plot_model: How do I write a new block and give it a name

Issue

Is there a possibility to create my own block using keras

New_Block

Solution

Maybe something like this:

import tensorflow as tf

class NewLayer(tf.keras.layers.Layer):
    def __init__(self):
      super(NewLayer, self).__init__()
      self.dense1 = tf.keras.layers.Dense(64, activation='relu')
      self.dense2 = tf.keras.layers.Dense(16, activation='relu')
      self.dense3 = tf.keras.layers.Dense(3)
      self._name = 'My_model'


    def call(self, inputs):
        x = tf.keras.layers.Flatten()(inputs)
        x = self.dense1(x)
        x = self.dense2(x)
        output = self.dense3(x)
        return output

inputs = tf.keras.layers.Input((128, 128, 3))
x = tf.keras.layers.Conv2D(256, 67)(inputs)
new_layer = NewLayer()
outputs = new_layer(x)
model = tf.keras.Model(inputs, outputs)

dot_img_file = 'model.png'
tf.keras.utils.plot_model(model, to_file=dot_img_file, show_shapes=True)

enter image description here

Answered By – AloneTogether

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published