TensorFlow自定义层拼接张量最后维度报错，求可行解决方法

你的问题核心是拼接的两个张量维度不匹配：输入张量是3维（batch_size, 10, 3），但位置嵌入数组是2维（10,1），缺少batch维度，导致TensorFlow无法对齐拼接。只需要给位置嵌入补上batch维度，并让它和输入的batch大小一致即可。

修改后的自定义层代码 #

直接在call方法中处理位置嵌入的维度扩展，确保和输入张量维度对齐：

import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import Input, Dense, Flatten

class PositionEmbeddingConcat(tf.keras.layers.Layer):
    def __init__(self, sequence_length, **kwargs):
        super(PositionEmbeddingConcat, self).__init__(**kwargs)
        # 初始化位置嵌入数组（保持原有逻辑）
        self.positional_embeddings_array = np.arange(sequence_length).reshape(sequence_length, 1)
        
    def call(self, inputs):
        # 将numpy数组转为TensorFlow张量，匹配输入数据类型
        pos_emb = tf.convert_to_tensor(self.positional_embeddings_array, dtype=inputs.dtype)
        # 扩展batch维度：从(10,1)变为(1,10,1)
        pos_emb = tf.expand_dims(pos_emb, axis=0)
        # 广播匹配输入的batch大小，自动适配任意batch_size
        pos_emb = tf.broadcast_to(pos_emb, tf.shape(inputs)[:2] + (1,))
        # 按最后一维拼接
        outp = tf.concat([inputs, pos_emb], axis=2)
        return outp

seq_len = 10    

input_layer = Input(shape=(seq_len, 3))
embedding_layer = PositionEmbeddingConcat(sequence_length=seq_len)
embeddings = embedding_layer(input_layer)
dense_layer = Dense(units=1)
output = dense_layer(Flatten()(embeddings))
modelT = tf.keras.Model(input_layer, output)

# 验证形状
print(modelT.output_shape)  # 输出 (None, 1)，符合预期

关键改动说明 #

1. 维度扩展：用tf.expand_dims给位置嵌入增加batch维度（从2维变3维）
2. 动态匹配batch大小：用tf.broadcast_to让位置嵌入的batch维度自动适配输入的batch_size，不管输入是多大的批次都能兼容
3. 类型对齐：将numpy数组转为TensorFlow张量时，指定和输入一致的数据类型，避免类型不匹配错误

另外，也可以用tf.tile替代tf.broadcast_to，效果相同：

batch_size = tf.shape(inputs)[0]
pos_emb = tf.tile(pos_emb, [batch_size, 1, 1])

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