Keras自定义Lambda层仅调用一次却重复执行两次的问题求解

执行两次的根本原因 #

这是TensorFlow/Keras的正常计算图构建机制导致的：

• 第一次执行发生在计算图构建阶段：此时输入是无实际数值的占位符张量，框架需要运行一次你的mpsm_process函数推导张量形状、生成运算节点，完成图构建
• 第二次执行发生在实际运行阶段（训练/预测）：此时输入带真实数值，逻辑执行一次输出实际计算结果

现有代码的隐藏问题 #

你当前的实现存在两个会导致后续运行报错的风险点：

• 硬编码固定batch_size=16：实际运行时batch size可能和设定值不一致，会直接触发形状不匹配报错，需要改为动态获取batch维度
• Lambda层内部循环创建Flatten()层实例：每次循环都会生成新的层实例，会导致计算图结构混乱，应该直接使用纯张量运算函数实现reshape逻辑

具体修复方法 #

1. 解决打印重复问题 #

将代码中所有Python原生print替换为TensorFlow内置的tf.print，tf.print是计算图运算节点，只会在图实际运行时输出内容，不会在构图阶段打印，可直接解决重复输出的问题。

2. 修正硬编码和层实例问题 #

修改mpsm_process函数的核心逻辑，同时替换低效的Python循环为向量化运算，大幅提升运行效率：

def mpsm_process(rfam_output):
    # 动态获取batch size，不要硬编码
    batch_size = tf.shape(rfam_output[0])[0]

    i_low = rfam_output[0]
    f_low = rfam_output[1]
    i_mid = rfam_output[2]
    f_mid = rfam_output[3]
    i_high = rfam_output[4]
    f_high = rfam_output[5]

    U_low = i_low + f_low
    tf.print("low shape: ", tf.shape(U_low))
    U_mid = i_mid + f_mid
    tf.print("mid shape: ", tf.shape(U_mid))
    U_high = i_high + f_high
    tf.print("high shape: ", tf.shape(U_high))

    # 动态维度用tf.shape获取，静态维度可以用.shape
    H_low = tf.shape(U_low)[1]
    W_low = tf.shape(U_low)[2]
    H_mid = tf.shape(U_mid)[1]
    W_mid = tf.shape(U_mid)[2]
    H_high = tf.shape(U_high)[1]
    W_high = tf.shape(U_high)[2]

    U_low = U_low[:, ::(H_low//H_high), ::(W_low//W_high), :]
    U_mid = U_mid[:, ::(H_mid//H_high), ::(W_mid//W_high), :]

    U_low = U_low[:, :H_high, :W_high, :]
    U_mid = U_mid[:, :H_high, :W_high, :]

    U_concat = concatenate([U_low, U_mid, U_high], name='U_concat')

    H_U = tf.shape(U_concat)[1]
    W_U = tf.shape(U_concat)[2]

    U_concat = U_concat[:, :(H_U - H_U % 5), :(W_U - W_U % 5), :]
    tf.print("concat shape: ", tf.shape(U_concat))

    cube_size = H_U // 5
    # 内置API切分patch，替换Python循环
    U_patches = tf.image.extract_patches(
        images=U_concat,
        sizes=[1, cube_size, cube_size, 1],
        strides=[1, cube_size, cube_size, 1],
        rates=[1,1,1,1],
        padding='VALID'
    )
    # U_patches形状调整为(batch, 25, 特征维度)
    U_patches = tf.reshape(U_patches, (batch_size, 25, -1))

    # 向量化计算两两余弦相似度，替换4层Python循环
    norm_patches = tf.nn.l2_normalize(U_patches, axis=-1)
    similarity_map = tf.matmul(norm_patches, norm_patches, transpose_b=True)
    # 扩展维度匹配原来的输出形状要求
    similarity = tf.expand_dims(similarity_map, axis=-1)
    tf.print(tf.shape(similarity))

    return similarity

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