tensorflowautoencoder

Accessing latent space of submodel of larger model


I have this autoencoder class

class SimpleAE(tf.keras.Model):
def __init__(self, latent_dim, bypass = False, trainable=True, **kwargs):
    super(SimpleAE,self).__init__(**kwargs)
    self.latent_dim = latent_dim
    self.bypass = bypass
    self.trainable = trainable
    self.quantizer = None
    
def get_config(self):
    config = super(SimpleAE,self).get_config().copy()
    config.update({'latent_dim':self.latent_dim, 'bypass':self.bypass, 'trainable':self.trainable, 
                   'quantizer':self.quantizer})        
    
    return config


    
    

    
def build(self,input_shape):   
    self.inputlayer = tf.keras.layers.InputLayer(input_shape=(input_shape[-1],))  # Initialize input layer
    self.encoder = tf.keras.layers.Dense(self.latent_dim,activation='linear', name="latentspace")
    self.decoder = tf.keras.layers.Dense(input_shape[-1], activation='linear')
    self.built = True
    
def call(self,x):
    if not self.built:
        self.build(x.shape)  # Ensure the model is built before calling
    # print("Test",flush=True)
    if self.bypass is False:
        xin = self.inputlayer(x)
        encoded = self.encoder(xin)
        decoded = self.decoder(encoded)
        return decoded
    else:

        return x

which is used in another model:

class Model():
def __init__(self, args, pca = None):   
    self.N = args.N
    self.L = args.L
    self.B = args.B
    self.H = args.H
    self.S = args.S
    self.P = args.P
    self.X = args.X
    self.R = args.R
    self.GPU = True #args.GPU
    self.skip = args.skip
    self.top = args.topology
    self.causal = args.causal
    self.duration = args.duration
    self.sample_rate = args.sample_rate
    self.model_name = self.top + "_auto_encoder"
    self.buffer_length = 10 #args.buffer_length
    self.save_latent = args.save_latent
    # self.pca = pca
    self.quantizer = None
             
    ...SOMEMODELSTUFF...
    
    self.AE_enc_left = SimpleAE(latent_dim=args.ae_enc_dim, bypass=args.bypass, trainable=True, name = "AE_Encoder_left")
    self.AE_enc_right = SimpleAE(latent_dim=args.ae_enc_dim, bypass=args.bypass, trainable=True, name = "AE_Encoder_right")
    self.AE_tcn_left = SimpleAE(latent_dim=args.ae_tcn_dim, bypass=args.bypass, trainable=True, name = "AE_TCN_left")
    self.AE_tcn_right = SimpleAE(latent_dim=args.ae_tcn_dim, bypass=args.bypass, trainable=True, name = "AE_TCN_right")
    
    

  
    

def call(self):

        input_left  = tf.keras.Input(shape = (None,), name = "Input_left")
        ... MOREMODELSTUFF ...

        

            
        enc_inp_r_quantized = self.AE_enc_right(output_of_submodel)
        enc_inp_l_quantized = self.AE_enc_left(output_of_submodel)
        ... FURTER  MODEL STUFF ...

Now I pretrained the entire model including the SimpleAutoencoders and load the model:

    model.load_weights(weights_file)
model.compile(loss=loss)

#sanity check
model.evaluate(valid_ds, verbose=0)
submodel = tf.keras.Model(inputs=[model.input], outputs=[model.get_layer("AE_Encoder_left").get_layer("latentspace").output])
print("test")

model.evaluate yields the expected output, so everything should be loaded correctly. "bypass" is set to false for the autoencoder, I checked by printing within the call method. When I look at the Autoencoder weights, they look normal, all weights seem to be there. HOWEVER,

model.get_layer("AE_Encoder_left").get_layer("latentspace").output always yields the following error *** AttributeError: Layer latentspace has no inbound nodes.

However, I want to continue by generating the latent space data from the autoencoders to train some quantizer. How to I get the output of the latentspace layer of the SimpleAE class as a part of the larger model?


Solution

  • I searched alot on the internet and found two somewhat opposing statements regarding this issue. One is referring to this behavior as a bug, the other as intended behavior, nontheless, apparently you cannot do it. One work around would be to store the output in a variable in the call method. However, I chose a different route: I split up my autoencoder into two seperate models, the encoder and the decoder. Then I can access the output of the encoder, which is the latent space.