如何将2D tensor按指定维度堆叠为3D tensor，得到目标形状[32,512,1536]

错误原因分析 #

你之前得到错误形状的核心问题有两个：

• 仅对2D张量在axis=1扩维后得到的形状为[32, 1, 768]，未将该维度扩展到和3D张量一致的512长度
• 拼接时选错了轴，在axis=1拼接会累加该维度的长度，得到[32, 512+1=513, 768]的错误形状

正确操作逻辑 #

要得到目标形状[32,512,1536]，需要先对齐两个张量的前两个维度，再在最后一个特征维度拼接：

1. 对形状为[32,768]的2D张量，先在axis=1插入新维度得到[32, 1, 768]，再将该维度复制512次，得到形状为[32, 512, 768]的3D张量
2. 将处理后的2D张量与原始3D张量在axis=-1（最后一个维度）拼接，两个768长度的特征维度相加得到1536，最终输出目标形状

不同框架代码示例 #

PyTorch实现 #

import torch

# 示例张量定义
tensor_2d = torch.randn(32, 768)
tensor_3d = torch.randn(32, 512, 768)

# 处理2D张量，用expand生成视图更省内存，无需实际复制数据
tensor_2d_expanded = tensor_2d.unsqueeze(1).expand(-1, 512, -1)
# 沿最后一个维度拼接
result = torch.cat([tensor_3d, tensor_2d_expanded], dim=-1)

print(result.shape) # 输出：torch.Size([32, 512, 1536])

TensorFlow实现 #

import tensorflow as tf

# 示例张量定义
tensor_2d = tf.random.normal((32, 768))
tensor_3d = tf.random.normal((32, 512, 768))

# 处理2D张量
tensor_2d_expanded = tf.repeat(tf.expand_dims(tensor_2d, axis=1), repeats=512, axis=1)
# 沿最后一个维度拼接
result = tf.concat([tensor_3d, tensor_2d_expanded], axis=-1)

print(result.shape) # 输出：(32, 512, 1536)

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