Recurrent loss in tensorflow without executing eagerly

Question

I have the following very simple example for a loss (that probably doesnt make sense)

import tensorflow as tf

class Loss:
  def __init__(self):
    self.last_output = tf.constant([0.5,0.5])

  def recurrent_loss(self, model_output):
    now = 0.9*self.last_output + 0.1*model_output
    self.last_output = now
    return tf.reduce_mean(now)

that just evaluates the reduced_mean of the model_output combined with the last model_output (in a ratio of 9 to 1). So for example

>> l = Loss()
>> l.recurrent_loss(tf.constant([1.,1.]))
tf.Tensor(0.55, shape=(), dtype=float32)
>> l.recurrent_loss(tf.constant([1.,1.]))
tf.Tensor(0.595, shape=(), dtype=float32)

if I understand correctly how tf is working this is only possible because by default tf is executing eagerly (tf.executing_eagerly() == True). That's should be the reason why I can overwrite the self.last_output variable with a new tensor to achieve a recurrent structure.

My Question: How can I can i achieve the same kind of recurrent structure in a tf graph that doesn't use eager execution?

Answer 1

in graph mode you have to use tf.Variable that is only created the first time the function is executed e.g :

class Loss:
  def __init__(self):
    self.last_output = None
  @tf.function
  def recurrent_loss(self, model_output):
    if  self.last_output is None:
        
        self.last_output = tf.Variable([0.5,0.5])
        
    now = 0.9*self.last_output + 0.1*model_output
    self.last_output.assign(now)
    return tf.reduce_mean(now)