如何通过均值/中位数插补统一张量长度并实现Keras层点积合并？

你的核心问题是批量输入下，tf.reduce_mean()返回的是每个样本的均值（形状(batch_size,)），而非单个标量，直接填充会导致形状不匹配，进而触发Dot层的报错。下面是修正后的实现方案，完全贴合你的需求：

修正后的自定义层代码 #

import tensorflow as tf
from tensorflow.keras.layers import Layer, Dot

class MeanPadDot(Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.dot_layer = Dot(axes=1)
    
    def call(self, inputs):
        x1, x2 = inputs
        # 获取两个张量的特征维度长度（动态形状，适配NAS的随机单元数）
        len1 = tf.shape(x1)[1]
        len2 = tf.shape(x2)[1]
        
        # 区分长、短张量并计算长度差
        long_tensor, short_tensor = (x1, x2) if len1 > len2 else (x2, x1)
        diff = tf.abs(len1 - len2)
        
        # 计算每个样本的均值，形状为(batch_size,)
        mean_vals = tf.reduce_mean(short_tensor, axis=1)
        # 将均值扩展为(batch_size, diff)：先增加维度再复制diff次
        mean_pad = tf.tile(tf.expand_dims(mean_vals, axis=1), [1, diff])
        # 拼接短张量与填充张量，得到和长张量同长度的张量
        padded_short = tf.concat([short_tensor, mean_pad], axis=1)
        
        # 计算点积
        return self.dot_layer([long_tensor, padded_short])
    
    def compute_output_shape(self, input_shape):
        # 点积输出形状为(batch_size,)
        return (input_shape[0][0],)

测试代码 #

# 定义输入层，适配任意长度的1D张量
input1 = tf.keras.Input(shape=(None,))
input2 = tf.keras.Input(shape=(None,))

# 实例化自定义层并构建模型
merge_layer = MeanPadDot()([input1, input2])
model = tf.keras.Model(inputs=[input1, input2], outputs=merge_layer)

# 生成测试用的批量输入（两个张量长度不同）
test_batch1 = tf.random.normal((4, 7))  # 批量4，长度7
test_batch2 = tf.random.normal((4, 4))  # 批量4，长度4

# 验证输出
output = model.predict([test_batch1, test_batch2])
print("输出形状:", output.shape)  # 应为(4,)
print("输出结果:", output)

关键修正说明 #

1. 动态形状处理：用tf.shape()获取张量的运行时形状，适配NAS生成的随机Dense单元数，避免静态形状不兼容的问题。
2. 批量均值填充：通过tf.expand_dims()把每个样本的均值从(batch_size,)扩展为(batch_size, 1)，再用tf.tile()复制diff次，生成和短张量同批量的填充张量，保证每个样本的填充值是自身的均值。
3. 形状对齐：拼接后padded_short与long_tensor的形状完全一致（(batch_size, max_len)），此时Dot层可以正常计算点积。

如果要替换成中位数填充，只需要把tf.reduce_mean换成tf.stats.reduce_median即可，逻辑完全通用。

内容的提问来源于stack exchange，提问作者David W. Thrower