Is there way to get the current learning rate or current epoch/step from within a custom tensorflow layer?

Luca Nunziante 0
Viewed 27

I know that it is possible to get the current learning rate simply by doing self.optimizer.lr when you are in your custom model, but I need to do something similar when implementing my own layer.
For now I have solved the issue by creating a function in my custom layer that accepts it as a parameter and is called by my custom model, but I was wondering if there is another way since there are many layers in my architecture and this way is pretty awful to see. I leave some code to be clearer.
For my purpose, it would be enough even just to get the current epoch or current step from inside the layer.
I am working in tensorflow 2.8.2. Thank you.
My layer structure is as follows

class my_layer(tf.keras.layers.Layer):
 #constructor, build, call methods etc..
 def function_to_get_lr(self,lr):
  #do sth with the lr

and in my custom model I do something like this

class my_model(tf.keras.Model):
 #other functions, constructor etc..
 def call(self, inputs, training=False):
  if training:
   for layer in self.layers:
    if "my_layer" in layer.name:
     layer.function_to_get_lr(self.optimizer.lr)
keras-layer tensorflow
1 Answers

I'm pretty sure that there is no "nice way" to do this, but you can do something like this:

def CustomLayer(Layer):
  def __init__(optimizer, ...):
    self.optimizer = optimizer
  def call(...):
    lr = self.optimizer.lr

class M(Model):
  def build(...):
    this.custom_layer = CustomLayer(self.optimizer, ...)
    ...
Alberto Sinigaglia0
Related