How to write class for neural network to be able to access layers in Tensorflow

Question

I have been training a network written in Tensorflow and all was well, but now I need to visualise some metrics related to the weights of the filters at each each layer. So I need to be able to access the layers, I know it's pretty simple when using Keras' Model , but I believe there is a better way to define the class using Tensorflow such that layers are accessible?

My code:

import tensorflow as tf
from tensorflow.keras.layers import (Activation, Conv2D, BatchNormalization)
from tensorflow.keras.optimizers import Adam

class neural_net():

    def __init__(self, x):

        self.x = x
        self.layers = self.predictor(self.x)
        self.loss = tf.reduce_mean(tf.squared_difference(self.layers[3], self.x), name = 'mse_loss')
    
    def predictor(self, x):
        layer0 = Conv2D(64, kernel_size=(3, 3), input_shape=(512,512,3), padding='same', name = 'conv1')(x)
        layer1 = Activation('relu')(layer0)
        layer2 = BatchNormalization()(layer1)
        pred = Conv2D(3, (3,3), activation='sigmoid', padding='same',name = 'decoded')(layer2)
        return [layer0,layer1,layer2,pred]

x = tf.compat.v1.placeholder(tf.float32, shape=(None,512,512,3),name = 'x')
net = neural_net(x)
loss = net.loss
layers = net.layers
train_op = tf.train.AdamOptimizer(1e-4).minimize(loss)

w=[]
for layer in net.layers():
    w.append(layer.get_weights())

This gives error:

AttributeError: 'Tensor' object has no attribute 'get_weights'

I have also seen some examples showing class with build and call methods but I couldn't figure how to use them for my case (limited knowledge of oop).

Answer 1

I would create a list like you suggested and put the layers into it. Then you can modify your predictor method to iterate the layers and call them in the right order. Example:

import tensorflow as tf
from tensorflow.keras.layers import (Activation, Conv2D, BatchNormalization)
from tensorflow.keras.optimizers import Adam

class neural_net():

    def __init__(self, x, y):

        self.x = x
        self.layers = [Conv2D(64, kernel_size=(3, 3), input_shape=(512,512,3), padding='same', name = 'conv1'), 
                       Activation('relu'), 
                       BatchNormalization(), 
                       Conv2D(3, (3,3), activation='sigmoid', padding='same',name = 'decoded')]
        self.pred = self.predictor(self.x)
        self.loss = tf.reduce_mean(tf.squared_difference(self.pred, self.x), name = 'mse_loss')
    
    def predictor(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

x = tf.compat.v1.placeholder(tf.float32, shape=(None,512,512,3),name = 'x')
net = neural_net(x)
loss = net.loss
pred = net.pred
train_op = tf.train.AdamOptimizer(1e-4).minimize(loss)

While I think this is what you meant, I doubt this would work like you expect it to. What is the exact reason you do not want to use the api provided by Keras? You can get the weights by iterating over the layers as well like you said.

Answer 2

I think what you should do is use the Keras Subclassing API (if you can call it an API). Essentially, what it does is inherit from the Keras Model (a little like the Sequential API), and so you can use it the same. It will enable your model to emulate the behaviors of a Keras model, for instance, you'll be able to call model.fit() or model.summary() .

Then, you should define the layers in the constructor (the __init__ part), which will set the layers as instance variables, and allow you to access the layers like model.layer1 . The call() method you will need to define will determine it's behavior when the model is called (the tensors will pass through the layers). Other than that, everything will be done automatically because the behaviors have been pre-defined by the Model class.

Here's an example on how to do it:

import tensorflow as tf
from tensorflow.keras import Model

inputs = tf.random.uniform(shape=(8, 512, 512, 3), minval=0, maxval=1)


class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer0 = tf.keras.layers.Conv2D(64, kernel_size=(3, 3), 
            padding='same', name='conv1')
        self.layer1 = tf.keras.layers.Activation('relu')
        self.layer2 = tf.keras.layers.BatchNormalization()

    def call(self, x, **kwargs):
        x = self.layer0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        return x


model = MyModel()

model.build(input_shape=(1, 512, 512, 3))

model(inputs)

You will also be able to access the individual layers like you're used to:

model.layers

[<tensorflow.python.keras.layers.convolutional.Conv2D at 0x2d337c4b9c8>,
 <tensorflow.python.keras.layers.core.Activation at 0x2d337c460c8>,
 <tensorflow.python.keras.layers.normalization_v2.BatchNormalization at 0x2d337c465c8>]

If you want to integrate the loss and predict function to the model itself, I'm afraid it will require some level of skill with OOP.

Let me know if anything is missing from my answer. You can find similar Subclassing API examples with toy datasets in one of my repos , namely with the MNIST and Iris dataset.