如何获取大模型中子模型SimpleAE的latent space输出？

我定义了如下SimpleAE自编码器类：

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

该类被嵌入到另一大模型中：

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 ...

我已预训练包含SimpleAE的整个大模型并加载权重：

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输出符合预期，权重加载正常，且autoencoder的bypass已设为False。但执行model.get_layer("AE_Encoder_left").get_layer("latentspace").output时，始终报错：

*** AttributeError: Layer latentspace has no inbound nodes.

我需要获取该autoencoder的latent space数据以训练量化器，请问如何正确获取大模型中SimpleAE类latentspace层的输出？

问题根源 #

报错的核心原因是SimpleAE内部的latentspace层（即self.encoder）没有被Keras的模型计算图正确追踪——你在call方法里手动创建了InputLayer并连接，但这种方式没有将子层的连接关系注册到大模型的计算图中，导致Keras无法识别latentspace层的输入节点。

可行修正方案 #

方案1：修改SimpleAE实现，让子层自动注册到模型图 #

去掉call方法里的InputLayer，直接用输入张量连接各层，让Keras自动追踪层间关系：

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
        # 提前定义encoder，decoder在build时根据输入形状初始化
        self.encoder = tf.keras.layers.Dense(self.latent_dim,activation='linear', name="latentspace")
        self.decoder = 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.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)
        if self.bypass is False:
            encoded = self.encoder(x)  # 直接用输入张量连接encoder，无需额外InputLayer
            decoded = self.decoder(encoded)
            return decoded
        else:
            return x

修改后重新训练或加载权重，此时latentspace层的连接关系会被正确注册到大模型的计算图中，你原有的获取子层输出的代码即可正常运行。

方案2：不修改SimpleAE，构建大模型时直接记录中间输出 #

如果不想改动SimpleAE的代码，可以在大模型的call方法中，直接提取AE_enc_left的encoder输出张量并作为模型的额外输出：

def call(self):
    input_left  = tf.keras.Input(shape = (None,), name = "Input_left")
    ... MOREMODELSTUFF ...
    
    # 调用AE_enc_left时，拆分encoder和decoder的调用，记录latent输出
    ae_left = self.AE_enc_left
    encoded_left = ae_left.encoder(output_of_submodel)
    enc_inp_l_quantized = ae_left.decoder(encoded_left)
    
    # 同理处理AE_enc_right
    ae_right = self.AE_enc_right
    encoded_right = ae_right.encoder(output_of_submodel)
    enc_inp_r_quantized = ae_right.decoder(encoded_right)
    
    ... FURTER  MODEL STUFF ...
    
    # 将主输出和latent输出一起返回
    return [main_output, encoded_left]  # main_output为原模型的主输出

加载模型后，直接创建子模型提取对应的latent输出：

submodel = tf.keras.Model(inputs=[model.input], outputs=[model.output[1]])

方案3：用Keras后端函数直接计算latent输出 #

如果上述方案都无法实施，可以用tf.keras.backend.function直接构建输入到latent输出的计算函数：

# 获取目标层
ae_layer = model.get_layer("AE_Encoder_left")
latent_layer = ae_layer.get_layer("latentspace")

# 创建计算函数
get_latent = tf.keras.backend.function(
    inputs=[model.input],
    outputs=[latent_layer(ae_layer.input)]
)

# 用验证集数据计算latent空间数据
latent_data = get_latent(valid_ds)

内容的提问来源于stack exchange，提问作者user20780