如何将形状为[5,2,18,4096]的张量堆叠为[2,90,4096]？

假设你使用PyTorch或TensorFlow框架，以下是两种高效实现方式：

方法一：直接拼接子张量 #

这种方式直观易懂，先提取第0维度的所有子张量，再沿目标维度拼接：

PyTorch 实现 #

import torch

# 构造示例张量
x = torch.randn(5, 2, 18, 4096)
# 提取第0维度的所有子张量，沿dim=1拼接（对应子张量的第1维度，拼接后扩展为90）
result = torch.cat([x[i] for i in range(x.size(0))], dim=1)
print(result.shape)  # 输出: torch.Size([2, 90, 4096])

TensorFlow 实现 #

import tensorflow as tf

# 构造示例张量
x = tf.random.normal((5, 2, 18, 4096))
# 提取第0维度的所有子张量，沿axis=1拼接
result = tf.concat([x[i] for i in range(tf.shape(x)[0])], axis=1)
print(result.shape)  # 输出: (2, 90, 4096)

方法二：维度转置+重塑（更高效） #

通过转置调整维度顺序，再合并目标维度，避免循环提取子张量，性能更优：

PyTorch 实现 #

import torch

x = torch.randn(5, 2, 18, 4096)
# 转置维度：将原维度(5,2,18,4096)调整为(2,5,18,4096)
x_transposed = x.permute(1, 0, 2, 3)
# 合并第1、2维度（5和18），得到(2, 5*18, 4096)
result = x_transposed.reshape(2, -1, 4096)
# 或使用flatten简化：result = x.permute(1,0,2,3).flatten(1,2)
print(result.shape)  # 输出: torch.Size([2, 90, 4096])

TensorFlow 实现 #

import tensorflow as tf

x = tf.random.normal((5, 2, 18, 4096))
# 转置维度：将原维度(5,2,18,4096)调整为(2,5,18,4096)
x_transposed = tf.transpose(x, perm=[1, 0, 2, 3])
# 合并第1、2维度
result = tf.reshape(x_transposed, (2, -1, 4096))
print(result.shape)  # 输出: (2, 90, 4096)

关键逻辑说明 #

原张量形状为[5,2,18,4096]，第0维度的每个子张量形状是[2,18,4096]：

• 我们需要将5个这样的子张量在子张量的第1维度（对应原张量的第2维度）拼接，最终该维度长度变为18*5=90，得到目标形状[2,90,4096]。
• 转置+重塑的方式无需循环，更适合处理大规模张量。

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