Speeding up model with feedback to the encoder

TensorFlow General Discussion
1

I got the following model, an autoencoder with feedback to the input layers of both, encoder and decoder, which is very slow at the moment due to the for loop. But it seems TOO slow, even for that case. Is it possible to speed up the inference/training?
The model is:

class FRAE(tf.keras.Model):     def __init__(self, latent_dim, shape, ht, n1, n2, n3, n4, bypass=False, trainable=True,**kwargs):         super(FRAE, self).__init__(**kwargs)         self.latent_dim = latent_dim         self.shape = shape         self.ht = ht         self.buffer = tf.Variable(initial_value=tf.zeros(shape=(1,shape[0] * self.ht), dtype=tf.float32))         self.bypass = bypass         self.quantizer = None         self.trainable = trainable                  self.l1 = tf.keras.layers.Dense(n1, activation='swish', input_shape=shape)         self.l2 = tf.keras.layers.Dense(n1, activation='swish')         self.ls = tf.keras.layers.Dense(latent_dim, activation='swish')          self.l3 = tf.keras.layers.Dense(n3, activation='swish')         self.l4 = tf.keras.layers.Dense(n4, activation='swish')         self.l5 = tf.keras.layers.Dense(shape[-1], activation='linear')       def get_config(self):         config = super(FRAE,self).get_config().copy()         config.update({'latent_dim':self.latent_dim, 'bypass':self.bypass, 'quantizer':self.quantizer,                         "encoder":self.encoder, "buffer":self.buffer,                        'decoder':self.decoder,"ht":self.ht, "shape":self.shape, "name":self.name})                          return config                def update_buffer(self, new_element):         n = self.shape[0]         new_element_expanded = tf.expand_dims(new_element, axis=0)         self.buffer.assign(tf.keras.backend.concatenate([new_element_expanded, self.buffer[:, :-n]], axis=1))                  def resetBuffer(self):         self.buffer[:,:].assign(tf.zeros(shape=(1,self.shape[0] * self.ht), dtype=tf.float32))                  @tf.function     def call(self, x):                 if self.bypass is True:             print("Bypassing FRAE", flush=True)             return x         else:             x = tf.squeeze(x,axis=0)                          decoded = tf.TensorArray(tf.float32, size=tf.shape(x)[0])             for i in tf.range(tf.shape(x)[0]):                 xexpand = tf.expand_dims(x[i],axis=0)                 xin = tf.concat((xexpand, self.buffer), axis=1)                    encoded = self.ls(self.l2(self.l1(xin)))                 decin = tf.concat([encoded, self.buffer], axis=1)                 y = self.l5(self.l4(self.l3(decin)))                 decoded = decoded.write(i,y)                 i += 1                 self.update_buffer(tf.squeeze(y))                                              tmp = tf.transpose(decoded.stack(),[1,0,2])             return tmp
2

hi @Cola_Lightyear.
Did you note an improvement adding a @tf.function to the call method of your subclassing model?
Do you want to decorate as such your training function and maybe also your loss function, instead?

3

I did not see a great improvement using @tf.function, maybe slightly. I need a speedup by at least 2 orders of magnitude. Do you have an idea? Why would @tf.function on my loss and training function help? Do you mean like this btw:

@tf.function def MyTrainingFunction(): model =        setupmodel() model.fit(data,epochs=whatever)

?

I updated my current code in the OP btw

4

I really do not understand why this is this slow. It is similar to RNN code I saw, here it is just an autoencoder and the feedback occurs at two layers, that is all.