PyTorch中常用的张量维度操作方法解析
flatten：展平张量
在 PyTorch 中，flatten() 方法用于将多维张量压缩为一维形式。该方法会按行优先顺序将所有元素展开成单一维度的张量。这对于全连接层前的数据准备非常常见。
import torch
x = torch.tensor([[10, 20, 30],
[40, 50, 60]])
flattened = x.flatten()
print(flattened) # 输出: tensor([10, 20, 30, 40, 50, 60])
unsqueeze：扩展维度
unsqueeze(dim) 可在指定位置插入大小为 1 的新维度，常用于匹配模型输入形状或广播运算需求。维度索引从 0 开始计数。
# 原始二维张量
t = torch.rand(2, 5)
expanded = t.unsqueeze(2) # 在第3个维度增加一个轴
print(expanded.shape) # 输出: torch.Size([2, 5, 1])
squeeze：压缩单例维度
squeeze() 移除张量中所有尺寸为 1 的维度。若指定 dim 参数，则仅对该维度执行压缩（前提是该维度大小为 1）。
z = torch.zeros(1, 4, 1, 8)
squeezed_all = z.squeeze() # 移除所有 size=1 维度 → (4, 8)
squeezed_dim = z.squeeze(0) # 仅移除第0维 → (4, 1, 8)
print(squeezed_all.shape) # 输出: torch.Size([4, 8])
print(squeezed_dim.shape) # 输出: torch.Size([4, 1, 8])
view：重塑张量结构
view() 允许重新组织张量的形状，只要总元素数量保持不变。它不复制数据，而是返回原始数据的新"视图"，因此效率较高。
data = torch.arange(12)
reshaped = data.view(3, 4)
print(reshaped)
# 输出:
# tensor([[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
在 Transformer 架构中，常使用星号解包进行动态 reshape：
# head_shape = (batch_size, seq_len)
output = x.view(*head_shape, num_heads, head_dim)
flip：沿指定维度翻转元素
torch.flip(input, dims) 沿给定的一个或多个维度反转元素顺序。参数 dims 接受整数或整数列表。
mat = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
rev_rows = torch.flip(mat, dims=[0]) # 上下翻转
rev_both = torch.flip(mat, dims=[0,1]) # 同时上下左右翻转
print(rev_rows)
# tensor([[7, 8, 9],
# [4, 5, 6],
# [1, 2, 3]])
print(rev_both)
# tensor([[9, 8, 7],
# [6, 5, 4],
# [3, 2, 1]])
stack：沿新轴堆叠张量
torch.stack() 将多个相同形状的张量沿新增的维度拼接起来。结果张量的秩比原张量高一阶。
a = torch.tensor([1, 2])
b = torch.tensor([3, 4])
c = torch.tensor([5, 6])
stacked = torch.stack((a, b, c), dim=0)
print(stacked.shape) # 输出: torch.Size([3, 2])
print(stacked)
# tensor([[1, 2],
# [3, 4],
# [5, 6]])
chunk：分块切分张量
torch.chunk(tensor, chunks, dim) 将张量沿指定维度均分为若干块。若不能整除，则最后一块可能略小。
tensor_2d = torch.arange(18).reshape(3, 6)
# 按列分为两部分
parts = torch.chunk(tensor_2d, chunks=2, dim=1)
for idx, part in enumerate(parts):
print(f"子张量 {idx+1}:")
print(part)
# 子张量 1:
# tensor([[ 0, 1, 2],
# [ 6, 7, 8],
# [12, 13, 14]])
#
# 子张量 2:
# tensor([[ 3, 4, 5],
# [ 9, 10, 11],
# [15, 16, 17]])
