NumPy 与 PyTorch 维度操作对照

NumPy 和 PyTorch 的很多函数名称相似，但最容易出错的不是名字，而是输入怎样升维、在哪一轴拼接、索引张量必须是什么形状。这篇把原来两份几乎镜像的笔记合并起来。

快速对照

目的	NumPy	PyTorch
删除长度为 1 的轴	`np.squeeze`	`torch.squeeze`
插入新轴	`np.expand_dims`	`torch.unsqueeze`
在已有轴拼接	`np.concatenate`	`torch.cat`
新建一轴后堆叠	`np.stack`	`torch.stack`
水平拼接	`np.hstack`	`torch.hstack`
垂直拼接	`np.vstack`	`torch.vstack`
深度拼接	`np.dstack`	`torch.dstack`
将一维输入视为列	`np.column_stack`	`torch.column_stack`
沿轴按索引取值	`np.take_along_axis`	`torch.gather`

squeeze：只删除长度为 1 的轴

import numpy as np
import torch

a_np = np.zeros((1, 3, 1, 5))
a_t = torch.zeros(1, 3, 1, 5)

print(np.squeeze(a_np).shape)       # (3, 5)
print(torch.squeeze(a_t).shape)    # torch.Size([3, 5])

print(np.squeeze(a_np, axis=2).shape)   # (1, 3, 5)
print(torch.squeeze(a_t, dim=2).shape)  # torch.Size([1, 3, 5])

指定的轴长度不是 1 时，NumPy 会报错；PyTorch 的具体边界行为应查锁定版本。工程代码最好先断言预期形状，不让 batch 维因为 batch size 恰好为 1 而被无意删除。

expand_dims / unsqueeze：插入一条长度为 1 的轴

x_np = np.zeros((3, 5))
x_t = torch.zeros(3, 5)

print(np.expand_dims(x_np, axis=0).shape)  # (1, 3, 5)
print(torch.unsqueeze(x_t, dim=0).shape)   # (1, 3, 5)

print(np.expand_dims(x_np, axis=-1).shape) # (3, 5, 1)
print(torch.unsqueeze(x_t, dim=-1).shape)  # (3, 5, 1)

concatenate / cat：沿已有轴拼接

关键约束是：除拼接轴以外，其他轴的形状必须相同。

a_np = np.zeros((2, 3))
b_np = np.ones((4, 3))
out_np = np.concatenate([a_np, b_np], axis=0)

a_t = torch.zeros(2, 3)
b_t = torch.ones(4, 3)
out_t = torch.cat([a_t, b_t], dim=0)

print(out_np.shape, out_t.shape)  # (6, 3), (6, 3)

若沿轴 1 拼接，则轴 0 必须一致：

np.concatenate([np.zeros((2, 3)), np.ones((2, 4))], axis=1)
torch.cat([torch.zeros(2, 3), torch.ones(2, 4)], dim=1)

这个约束容易写反，这里按实际行为说明。

stack：先创建新轴，再堆叠

stack 要求所有输入形状相同。

items_np = [np.zeros((2, 3)), np.ones((2, 3))]
items_t = [torch.zeros(2, 3), torch.ones(2, 3)]

print(np.stack(items_np, axis=0).shape)    # (2, 2, 3)
print(torch.stack(items_t, dim=0).shape)   # (2, 2, 3)

print(np.stack(items_np, axis=1).shape)    # (2, 2, 3)
print(torch.stack(items_t, dim=1).shape)   # (2, 2, 3)

这里两个结果碰巧数字相同，但轴语义不同；真实代码应给每一轴命名或在注释中说明含义。

二维输入在最后一轴堆叠的示例如下：

a_np = np.arange(6).reshape(2, 3)
b_np = a_np + 10
a_t = torch.arange(6).reshape(2, 3)
b_t = a_t + 10

print(np.stack((a_np, b_np), axis=2).shape)   # (2, 3, 2)
print(torch.stack((a_t, b_t), dim=2).shape)  # torch.Size([2, 3, 2])

这里最后一维的两个值分别来自 a 和 b，而 concatenate/cat(..., axis=1/dim=1) 会得到 (2, 6)，不会新增轴。

hstack、vstack、dstack

这些便利函数会先按输入维数做升维，不能一概解释成固定的 axis=1。

hstack

一维输入：沿第一轴拼接；
二维及以上输入：沿第二轴拼接。

np.hstack([np.array([1, 2]), np.array([3, 4])])     # shape (4,)
torch.hstack([torch.tensor([1, 2]), torch.tensor([3, 4])])

vstack

一维输入先变成形如 (1, N) 的行，再沿第一轴拼接：

np.vstack([np.array([1, 2]), np.array([3, 4])])     # shape (2, 2)
torch.vstack([torch.tensor([1, 2]), torch.tensor([3, 4])])

dstack

一维输入先视为 (1, N, 1)，二维输入先视为 (M, N, 1)，再沿第三轴拼接：

np.dstack([np.array([1, 2]), np.array([3, 4])])     # shape (1, 2, 2)
torch.dstack([torch.tensor([1, 2]), torch.tensor([3, 4])])

把一维便利函数的例子放在一起看，输出最直观：

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

print(np.hstack((a, b)))
# [1 2 3 4 5 6]

print(np.vstack((a, b)))
# [[1 2 3]
#  [4 5 6]]

print(np.dstack((a, b)))
# [[[1 4]
#   [2 5]
#   [3 6]]]

将 np.array 换成 torch.tensor，三种 PyTorch 便利函数具有对应的维度提升规则；不要把一维 hstack 误写成固定的 dim=1。

二维 dstack 的完整输出是：

a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6], [7, 8]])
print(np.dstack((a, b)))
# [[[1 5]
#   [2 6]]
#  [[3 7]
#   [4 8]]]

column_stack：将一维输入作为列

left_np = np.array([1, 2, 3])
right_np = np.array([4, 5, 6])
print(np.column_stack([left_np, right_np]).shape)  # (3, 2)

left_t = torch.tensor([1, 2, 3])
right_t = torch.tensor([4, 5, 6])
print(torch.column_stack([left_t, right_t]).shape) # (3, 2)

二维输入不会再次转置，而是按列方向组合。

对应的数值输出与二维输入示例如下：

print(np.column_stack((left_np, right_np)))
# [[1 4]
#  [2 5]
#  [3 6]]

x = np.array([[1, 2], [3, 4]])
y = np.array([[5, 6], [7, 8]])
print(np.column_stack((x, y)))
# [[1 2 5 6]
#  [3 4 7 8]]

take_along_axis / gather

二者都表达“沿指定轴，用一个索引数组选择元素”，但形状和广播规则不是完全同一套 API。

NumPy：

scores = np.array([[0.2, 0.9, 0.1], [0.7, 0.3, 0.4]])
indices = np.array([[1], [0]])
chosen = np.take_along_axis(scores, indices, axis=1)
print(chosen)  # [[0.9], [0.7]]

PyTorch：

scores = torch.tensor([[0.2, 0.9, 0.1], [0.7, 0.3, 0.4]])
indices = torch.tensor([[1], [0]])
chosen = torch.gather(scores, dim=1, index=indices)
print(chosen)

对 torch.gather：

input 和 index 必须有相同的维数；
输出形状与 index 相同；
input 与 index 不会自动彼此广播；
索引 dtype 必须满足 API 要求，常见为 torch.long。

写强化学习代码时，建议在 gather 前打印或断言 Q、action index 和 batch/agent 轴，避免“能运行但取错轴”。

完整整数索引例子

NumPy 原例让每一行按自己的索引顺序取值：

arr = np.array([
    [10, 20, 30],
    [40, 50, 60],
    [70, 80, 90],
])
indices = np.array([
    [2, 1, 0],
    [0, 2, 1],
    [1, 0, 2],
])

result = np.take_along_axis(arr, indices, axis=1)
print(result)
# [[30 20 10]
#  [40 60 50]
#  [80 70 90]]

PyTorch 原例展示了相同形状约束下的 gather：

source = torch.tensor([[1, 2], [3, 4], [5, 6]])
index = torch.tensor([[0, 1], [1, 0], [0, 0]])

result = torch.gather(source, dim=1, index=index)
print(result)
# tensor([[1, 2],
#         [4, 3],
#         [5, 5]])

两个输出都与索引数组同形状，但这不表示两个 API 的所有广播规则相同；把代码从一个库移植到另一个库时仍应查当前版本文档。

随机数：使用显式生成器

NumPy 新代码优先使用 Generator：

rng = np.random.default_rng(2026)
uniform = rng.random((2, 3))
normal = rng.normal(loc=0.0, scale=1.0, size=(2, 3))
integers = rng.integers(low=0, high=10, size=(2, 3))
permutation = rng.permutation(10)

PyTorch 可使用独立 Generator：

generator = torch.Generator(device="cpu")
generator.manual_seed(2026)

uniform = torch.rand((2, 3), generator=generator)
normal = torch.randn((2, 3), generator=generator)
integers = torch.randint(0, 10, (2, 3), generator=generator)
permutation = torch.randperm(10, generator=generator)

固定种子不自动保证跨版本、跨设备和不同并行策略逐位一致。可复现实验还要记录库版本、设备、确定性选项和数据加载顺序。

常见分布与采样

NumPy Generator 覆盖 rand、randn、normal、randint、choice 和 uniform 的常见用途：

rng = np.random.default_rng(2026)

# [0, 1) 均匀分布
unit_uniform = rng.random((2, 3))

# [low, high) 均匀分布
bounded_uniform = rng.uniform(low=-1.0, high=1.0, size=(2, 3))

# 标准正态与自定义正态
standard_normal = rng.standard_normal((2, 3))
normal = rng.normal(loc=5.0, scale=2.0, size=(2, 3))

# 整数、随机排列和无放回抽样
integers = rng.integers(0, 10, size=(2, 3))
permutation = rng.permutation(10)
sample = rng.choice(np.array([1, 2, 3, 4, 5]), size=3, replace=False)

旧 np.random.rand(...)、randn(...) 等全局函数仍常见于历史代码，但独立 Generator 更便于隔离不同模块的随机状态。high 通常是不包含的上界；采样前应明确是否允许放回以及概率向量是否归一化。

PyTorch 对应写法：

generator = torch.Generator(device="cpu").manual_seed(2026)

unit_uniform = torch.rand((2, 3), generator=generator)
standard_normal = torch.randn((2, 3), generator=generator)
integers = torch.randint(0, 10, (2, 3), generator=generator)
permutation = torch.randperm(10, generator=generator)

# 自定义均值和标准差
normal = torch.normal(mean=5.0, std=2.0, size=(2, 3), generator=generator)

# 指定区间；uniform_ 是原地操作
bounded_uniform = torch.empty(2, 3).uniform_(-1.0, 1.0, generator=generator)

在 CUDA 上使用独立生成器时，生成器设备必须与目标张量设备约定一致。uniform_ 会修改已有张量；若该张量已参与计算图或被其他对象共享，应避免无意的原地修改。

历史代码常使用全局随机 API。阅读这类代码时，可以按下表映射到显式生成器：

旧 NumPy 写法	新代码中的 `Generator`	PyTorch 对应
`np.random.rand(d0, d1)`	`rng.random((d0, d1))`	`torch.rand((d0, d1), generator=g)`
`np.random.randn(d0, d1)`	`rng.standard_normal((d0, d1))`	`torch.randn((d0, d1), generator=g)`
`np.random.normal(mu, sigma, size)`	`rng.normal(mu, sigma, size)`	`torch.normal(mean=mu, std=sigma, size=size, generator=g)`
`np.random.randint(low, high, size)`	`rng.integers(low, high, size)`	`torch.randint(low, high, size, generator=g)`
`np.random.choice(values, size, replace)`	`rng.choice(values, size, replace=replace)`	常用 `randperm` 或 `multinomial`，语义需分别核对
`np.random.uniform(low, high, size)`	`rng.uniform(low, high, size)`	`torch.empty(size).uniform_(low, high, generator=g)`

旧教程打印的随机数只是某次示意输出；没有同时记录种子、版本和设备时，不应把那些具体小数当成可复现基准。

常用 API 形式与形状约束

下面按函数列出最常用的参数，同时避免把某个版本的完整签名复制成永久承诺：

# NumPy
np.squeeze(a, axis=None)
np.expand_dims(a, axis)
np.concatenate(arrays, axis=0)
np.stack(arrays, axis=0)
np.take_along_axis(arr, indices, axis)

# PyTorch
torch.squeeze(input, dim=None)
torch.unsqueeze(input, dim)
torch.cat(tensors, dim=0)
torch.stack(tensors, dim=0)
torch.gather(input, dim, index)

关键区别：

操作	输入约束	输出变化
`squeeze`	只能删除长度为 1 的轴	轴数减少
`expand_dims` / `unsqueeze`	插入位置必须在合法范围	新增长度为 1 的轴
`concatenate` / `cat`	除拼接轴外全部相同	拼接轴长度相加
`stack`	所有输入形状完全相同	新增一条轴
`take_along_axis` / `gather`	索引维度与所选轴满足各自 API 规则	形状主要由索引决定

失败示例比成功示例更重要

# 沿 axis=0 拼接时，axis=1 的长度 3 和 4 不一致，因此失败
np.concatenate([np.zeros((2, 3)), np.ones((4, 4))], axis=0)

# stack 不会自动补齐不同形状
torch.stack([torch.zeros(2, 3), torch.ones(4, 3)], dim=0)

生产代码不应依赖异常文本判断逻辑；在边界处主动检查：

assert all(item.ndim == items[0].ndim for item in items)
assert all(item.shape[1:] == items[0].shape[1:] for item in items)
batch = torch.cat(items, dim=0)

NumPy 与 PyTorch 间转换

array = np.arange(6, dtype=np.float32).reshape(2, 3)
tensor = torch.from_numpy(array)

CPU 上 from_numpy 可能与原数组共享内存，修改一方会影响另一方。需要独立副本时显式复制：

tensor = torch.from_numpy(array.copy())
array_again = tensor.detach().cpu().numpy().copy()

GPU 张量必须先 detach().cpu() 才能转 NumPy；带梯度的计算图也不应被悄悄跨库修改。

形状检查习惯

assert tensor.ndim == 3
assert tensor.shape[-1] == FEATURE_DIM

最有效的张量排错信息通常是：每一轴的语义、完整 shape、dtype、device，以及拼接/索引前后的预期。把这些写进测试，比记住更多函数别名更可靠。

官方参考

NumPy：squeeze、expand_dims、concatenate、stack、take_along_axis、随机数生成
PyTorch：squeeze、unsqueeze、cat、stack、gather、Generator