How to define a custom layer in subclassing?

General Discussion
, ,
#1

How to define a layer method in subclass while operating in keras?

There must be a self.layer = ... in the init? But what should it be? Like in pytorch there is self.layers = nn.ModuleList(layers)

I read about implementing tf.Module in defining layers on the link; but how to write self.layers = tf.Module ? so that self.layers can be used in a subclass later in the code!

class GTN(keras.Model): # layers.Layer keeps track of everything under the hood!
    
    def __init__(self, num_edge, num_channels, w_in, w_out, num_class,num_layers,norm):
        super(GTN, self).__init__()
        self.num_edge = num_edge
        self.num_channels = num_channels
        self.w_in = w_in
        self.w_out = w_out
        self.num_class = num_class
        self.num_layers = num_layers
        self.is_norm = norm
        
        layers = []

        for i in tf.range(num_layers):
            if i == 0:
                layers.append(GTLayer(num_edge, num_channels, first=True))
            else:
                layers.append(GTLayer(num_edge, num_channels, first=False))
        

def call(self, A, X, target_x, target):
A = tf.expand_dims(A, 0)
Ws = []

for i in range(self.num_layers):
    if i == 0:
        H, W = self.layers[i](A) #self.layers = nn.ModuleList(layers)
    else:
        H = self.normalization(H)
        H, W = self.layers[i](A, H)
    Ws.append(W)

File “/Users/anshumansinha/Desktop/GTN/model_tf_2.py”, line 108, in call H, W = self.layersi #self.layers = nn.ModuleList(layers) AttributeError: Exception encountered when calling layer “gtn” (type GTN).

‘GTN’ object has no attribute ‘layers’

#2

You should be subclassing tf.keras.layers.Layer instead. This will have an init method, a build method for defining states and a call method for computing those states.

1 Like
#4

I did the way you suggested! But then how will I be able to train the weights of this model (GTN) during training ? At the moment in the code below, I was able to write model.trainable_weights but when I implemented subclassing from tf.keras.layers.Layer then error message comes up as : Error >
grads = tape.gradient(loss, model.trainable_weights()) , TypeError: 'list' object is not callable

Code:

loss,y_train,Ws = model(A, node_features, train_node, train_target)

grads = tape.gradient(loss, model.trainable_weights())
            optimizer.apply_gradients(zip(grads, model.trainable_weights()))

You may like to view the code for GTN , it’s available in the below link : link to GTN (model)