pytorch 常用 tensor 操作

原创 python pytorch torch

1. Tensor 的基本信息

Tensor 可以理解为 PyTorch 中的多维数组。每个 Tensor 都有形状、数据类型和所在设备等属性，下面是一个例子，其中 shape 或 size()：每个维度的大小；ndim 或 dim()：维度数量；numel() 或 nelement()：元素总数；dtype：元素的数据类型；device：Tensor 所在的设备，例如 cpu 或 cuda:0。

x = torch.tensor([[1, 2, 3], [4, 5, 6]])

print(x.shape, x.ndim, x.numel(), x.dtype, x.device)

'''
torch.Size([2, 3])
2
6
torch.int64
cpu
'''

2. 创建 Tensor

2.1 torch.tensor

torch.tensor() 可以直接根据 Python 列表创建 Tensor。只包含整数时通常推断为 torch.int64，包含 Python 浮点数时默认使用 PyTorch 当前的浮点类型，一般是 torch.float32。

integer_tensor = torch.tensor([[1, 2], [3, 4]])
float_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

print(integer_tensor.dtype)
print(float_tensor.dtype)

'''
torch.int64
torch.float32
'''

创建时可以通过 dtype 和 device 明确指定类型与设备，如果要修改设备类型，通常使用 to()、

x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, device='cpu',)
x = x.to(dtype=torch.float64)

if torch.cuda.is_available():
    x = x.to('cuda')

2.2 empty、zeros、ones 和 full

这几个函数可以直接创建 tensor，torch.empty() 只分配内存，不会把其中的值初始化为 0，因此刚创建时可能看到任意数值。需要确定初始值时，应使用 zeros、ones 或 full，表示创建全0，全1，或者指定数值的 tensor。

empty_tensor = torch.empty((2, 3))
zeros_tensor = torch.zeros((2, 3))
ones_tensor = torch.ones((2, 3))
full_tensor = torch.full((2, 3), 2.649)

print(zeros_tensor)
print(ones_tensor)
print(full_tensor)

2.3 随机 Tensor

常用的随机创建方法如下：

• torch.rand(size)：从 \([0, 1)\) 的均匀分布中采样；
• torch.randn(size)：从标准正态分布 \(\mathcal{N}(0,1)\) 中采样；
• torch.randint(low, high, size)：从 \([low, high)\) 中采样整数；
• torch.randperm(n)：返回 \([0,n)\) 的随机排列。

torch.manual_seed(0)
x1 = torch.rand((2, 3))
x2 = torch.randn((2, 3))
x3 = torch.randint(3, 10, (2, 3))
index = torch.randperm(4)

2.4 创建相同形状的 Tensor

*_like 系列函数会沿用输入 Tensor 的形状，并且默认沿用它的 dtype 和 device：

x = torch.tensor([[1, 2, 3], [4, 5, 6]])

zeros = torch.zeros_like(x)
ones = torch.ones_like(x)
random_values = torch.rand_like(x, dtype=torch.float32)

print(zeros)
print(ones)
print(random_values.shape)

类似的方法还有 torch.empty_like()、torch.randn_like()、torch.randint_like() 和 torch.full_like()。

3. dtype 和 FP16

3.1 常见数值类型

PyTorch 中常见的数据类型如下：

类型	PyTorch dtype	常见用途
16 位浮点数	torch.float16、torch.half	混合精度训练和推理
16 位脑浮点数	torch.bfloat16	更重视数值范围的混合精度计算
32 位浮点数	torch.float32、torch.float	默认浮点类型
64 位浮点数	torch.float64、torch.double	需要更高精度的数值计算
32 位整数	torch.int32、torch.int	普通整数数据
64 位整数	torch.int64、torch.long	默认整数类型，常用于索引和标签
布尔类型	torch.bool	掩码和逻辑判断

可以使用 to()、float()、half() 和 long() 等方法转换类型：

x = torch.tensor([1, 2, 3])

x_float = x.to(torch.float32)
x_half = x_float.half()
x_long = x_float.long()

3.2 FP16 的表示方法

FP16，也叫 binary16 或半精度浮点数，一共使用 16 位：

• 1 位符号位 \(S\)；
• 5 位指数位 \(E\)；
• 10 位尾数位 \(M\)。

当指数位在 00001 到 11110 之间时，它表示正规数：

\[ (-1)^S \times 2^{E-15} \times \left(1+\frac{M}{2^{10}}\right) \]

指数偏移量是 15。例如指数位的十进制值为 18 时，实际指数为：

\[ 18-15=3 \]

当指数位全为 0 时，如果尾数不为 0，则表示次正规数：

\[ (-1)^S \times 2^{-14} \times \frac{M}{2^{10}} \]

当指数位全为 1 时，尾数为 0 表示正负无穷大，尾数不为 0 表示 NaN。

FP16 的几个关键数值为：

数值	结果
最大有限正数	\(65504\)
最小正规正数	\(2^{-14}\approx6.10352\times10^{-5}\)
最小次正规正数	\(2^{-24}\approx5.96046\times10^{-8}\)

例如，FP16 位模式 0 01111 0000000000 中，符号位为 0、实际指数为 \(15-15=0\)、尾数为 1，因此它表示：

\[ 1\times2^0=1 \]

在 PyTorch 中可以使用 torch.finfo() 查看浮点类型的范围：

info = torch.finfo(torch.float16)

print(info.bits)
print(info.max)
print(info.smallest_normal)

'''
16
65504.0
6.103515625e-05
'''

FP16 可以减少显存占用并提高部分硬件上的计算速度，但有效位数较少，更容易出现舍入误差、上溢和下溢。因此训练中通常结合自动混合精度使用，而不是简单地把所有 Tensor 都强制转换为 FP16。

4. 调整形状和维度

4.1 view 和 reshape

view() 和 reshape() 都可以改变 Tensor 的形状，但不会改变元素总数：

x = torch.arange(24).reshape(2, 3, 4)

y1 = x.view(6, 4)
y2 = x.reshape(4, 6)

print(y1.shape)
print(y2.shape)

'''
torch.Size([6, 4])
torch.Size([4, 6])
'''

view() 要求新的形状与原来的 size 和 stride 兼容。reshape() 在条件允许时返回 View，否则会复制数据。无法确定内存是否连续时，使用 reshape() 更方便。

4.2 transpose 和 permute

transpose(dim0, dim1) 交换两个维度，permute(dims) 可以按照指定顺序重新排列所有维度：

x = torch.randn(2, 3, 4)

y1 = x.transpose(1, 2)
y2 = x.permute(2, 1, 0)

print(y1.shape)
print(y2.shape)

'''
torch.Size([2, 4, 3])
torch.Size([4, 3, 2])
'''

维度交换后得到的 Tensor 可能不再连续，此时直接使用 view() 可能报错：

x = torch.arange(24).reshape(2, 3, 4)
y = x.permute(2, 1, 0)

print(y.is_contiguous())

# 先转为连续内存，再使用 view
z1 = y.contiguous().view(4, 6)

# 或者直接使用 reshape
z2 = y.reshape(4, 6)

contiguous() 只会在需要时复制数据。更详细的 Storage、stride 和连续性原理可以单独参考博客中的 contiguous 文章。

4.3 squeeze 和 unsqueeze

squeeze() 用于删除大小为 1 的维度，unsqueeze() 用于在指定位置增加一个大小为 1 的维度：

x = torch.zeros(2, 1, 3)

y1 = x.squeeze(1)
y2 = y1.unsqueeze(1)

print(y1.shape)
print(y2.shape)

'''
torch.Size([2, 3])
torch.Size([2, 1, 3])
'''

不指定 dim 时，squeeze() 会删除所有大小为 1 的维度。处理 batch 数据时最好明确指定 dim，避免 batch size 为 1 时误删 batch 维度。

5. 扩展和重复

5.1 expand

Tensor.expand() 只能扩展大小为 1 的维度。参数中的 -1 表示保持该维度不变：

x = torch.tensor([[1], [2], [3]])

y1 = x.expand(3, 4)
y2 = x.expand(-1, 4)

print(y1)
print(y2.shape)

'''
tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3]])
torch.Size([3, 4])
'''

expand() 不会真正复制数据，而是通过把对应维度的 stride 设置为 0 创建 View。因此扩展后的多个元素可能指向同一块内存，不要直接对它执行 In-place Operation；确实需要写入时先调用 clone()。

5.2 repeat_interleave

repeat_interleave() 会逐个重复元素。不指定 dim 时，会先把输入展平：

x = torch.tensor([[1, 2], [3, 4]])

print(torch.repeat_interleave(x, 2))
print(torch.repeat_interleave(x, 3, dim=1))

'''
tensor([1, 1, 2, 2, 3, 3, 4, 4])
tensor([[1, 1, 1, 2, 2, 2],
        [3, 3, 3, 4, 4, 4]])
'''

也可以为不同元素指定不同的重复次数：

x = torch.tensor([[1, 2], [3, 4]])
repeats = torch.tensor([1, 2])

print(torch.repeat_interleave(x, repeats, dim=0))

'''
tensor([[1, 2],
        [3, 4],
        [3, 4]])
'''

6. 拼接和拆分

6.1 torch.cat

torch.cat(tensors, dim=0) 沿已有维度拼接多个 Tensor。除拼接维度外，其他维度的大小必须一致：

x = torch.tensor([[1, 2, 3], [4, 5, 6]])

row = torch.cat((x, x), dim=0)
column = torch.cat((x, x), dim=1)

print(row.shape)
print(column.shape)

'''
torch.Size([4, 3])
torch.Size([2, 6])
'''

6.2 torch.split

torch.split() 可以按照固定大小或指定列表拆分 Tensor，返回一个由多个 Tensor 组成的元组：

x = torch.arange(10).reshape(5, 2)

parts_1 = torch.split(x, 2, dim=0)
parts_2 = torch.split(x, [1, 4], dim=0)

print([part.shape for part in parts_1])
print([part.shape for part in parts_2])

'''
[torch.Size([2, 2]), torch.Size([2, 2]), torch.Size([1, 2])]
[torch.Size([1, 2]), torch.Size([4, 2])]
'''

当拆分大小不能整除该维度时，最后一块可以更小。

6.3 torch.chunk

torch.chunk(input, chunks, dim=0) 尝试把 Tensor 拆成指定数量的块：

x = torch.arange(10).reshape(2, 5)
parts = torch.chunk(x, 2, dim=1)

print(parts)

'''
(tensor([[0, 1, 2],
         [5, 6, 7]]),
 tensor([[3, 4],
         [8, 9]]))
'''

这里需要注意，torch.chunk() 在某些不能完成均匀拆分的情况下，返回的块数可能少于 chunks。如果必须得到准确数量的块，可以使用 torch.tensor_split()。

7. 统计和筛选

7.1 torch.max

不指定维度时，torch.max(input) 返回整个 Tensor 的最大值：

x = torch.tensor([[1.0, 5.0, 3.0], [4.0, 2.0, 6.0]])
print(torch.max(x))

'''
tensor(6.)
'''

指定 dim 后，会同时返回该维度上的最大值和索引：

values, indices = torch.max(x, dim=1)

print(values)
print(indices)

'''
tensor([5., 6.])
tensor([1, 2])
'''

keepdim=True 会保留被归约的维度，只把该维度的大小变成 1：

values, indices = torch.max(x, dim=1, keepdim=True)
print(values.shape)

'''
torch.Size([2, 1])
'''

7.2 torch.mean

torch.mean() 与 torch.max() 的 dim 和 keepdim 用法相似，但它计算的是均值：

x = torch.tensor([[1.0, 5.0, 3.0], [4.0, 2.0, 6.0]])

print(torch.mean(x))
print(torch.mean(x, dim=0))
print(torch.mean(x, dim=0, keepdim=True).shape)

'''
tensor(3.5000)
tensor([2.5000, 3.5000, 4.5000])
torch.Size([1, 3])
'''

输入必须是浮点或复数类型。整数 Tensor 需要先转换类型：

x = torch.tensor([1, 2, 3]).float()
print(x.mean())

7.3 torch.topk

torch.topk() 返回指定维度上最大的前 \(k\) 个值以及对应索引，分类任务中经常用它查找概率最高的类别：

scores = torch.tensor([
    [0.1, 0.8, 0.4, 0.6],
    [0.9, 0.2, 0.7, 0.3],
])

values, indices = torch.topk(scores, k=2, dim=1)

print(values)
print(indices)

'''
tensor([[0.8000, 0.6000],
        [0.9000, 0.7000]])
tensor([[1, 3],
        [0, 2]])
'''

设置 largest=False 可以返回最小的 \(k\) 个值，sorted=True 表示对返回结果进行排序。

8. 创建网格坐标

torch.meshgrid() 根据多个一维 Tensor 创建坐标网格。建议明确填写 indexing，避免依赖默认行为：

x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5, 6, 7])

grid_x, grid_y = torch.meshgrid(x, y, indexing='ij')

print(grid_x)
print(grid_y)

'''
tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3]])
tensor([[4, 5, 6, 7],
        [4, 5, 6, 7],
        [4, 5, 6, 7]])
'''

indexing='ij' 表示输出维度顺序与输入顺序一致。使用 indexing='xy' 时，前两个维度会按照笛卡尔坐标的习惯交换。

下面使用网格坐标绘制一个三维曲面：

import matplotlib.pyplot as plt

xs = torch.linspace(-5, 5, steps=100)
ys = torch.linspace(-5, 5, steps=100)
x, y = torch.meshgrid(xs, ys, indexing='xy')
z = torch.sin(torch.sqrt(x * x + y * y))

ax = plt.axes(projection='3d')
ax.plot_surface(x.numpy(), y.numpy(), z.numpy())
plt.show()

9. 使用 scatter 按索引写入

torch.scatter(input, dim, index, src) 会根据 index，沿 dim 指定的方向把 src 中的数据写入 input。

对于二维 Tensor，可以先记住下面两个关系：

self[index[i][j]][j] = src[i][j]  # dim == 0
self[i][index[i][j]] = src[i][j]  # dim == 1

下面以 dim=0 为例：

src = torch.arange(12, dtype=torch.float32).reshape(2, 6)
index = torch.tensor([
    [0, 1, 2, 0, 0],
    [2, 0, 0, 1, 2],
])
output = torch.zeros(3, 6)

result = torch.scatter(output, dim=0, index=index, src=src)
print(result)

'''
tensor([[ 0.,  7.,  8.,  3.,  4.,  0.],
        [ 0.,  1.,  0.,  9.,  0.,  0.],
        [ 6.,  0.,  2.,  0., 10.,  0.]])
'''

初始化时，src、index 和 output 的关系如下：

例如图中的序号 3：

目标行 = index[0][2] = 2
目标列 = 2
写入数值 = src[0][2] = 2

因此执行的是 output[2][2] = 2。

继续填充后，例如 index[1][4] = 2，对应操作为：

output[2][4] = src[1][4] = 10

这里 index 的第二维小于 src，因此 src 最后一列没有参与写入，对应的输出位置仍然保持初始值 0：

scatter_() 是它的原地版本：

output.scatter_(dim=0, index=index, src=src)

10. In-place Operation

In-place Operation 指直接修改原 Tensor 的操作，不另外返回一份独立结果。PyTorch 中带下划线后缀的方法通常是原地操作，例如：

x.add_(2)
x.squeeze_()
x.scatter_(dim, index, src)

+=、*= 等运算符也可能执行原地修改。可以通过 data_ptr() 观察底层数据地址：

x = torch.tensor([1.0, 2.0, 3.0])

before = x.data_ptr()
x.add_(2)
after_in_place = x.data_ptr()

y = x + 2
after_out_of_place = y.data_ptr()

print(before == after_in_place)
print(before == after_out_of_place)

'''
True
False
'''

原地操作有时可以减少额外分配，但在自动求导中需要谨慎。反向传播可能需要前向计算保存的中间值，如果这些值被原地修改，PyTorch 会检测到版本不一致并抛出错误。例如，不能直接原地修改一个需要梯度的叶子 Tensor：

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
x.add_(1)

# RuntimeError: a leaf Variable that requires grad is being used
# in an in-place operation

因此，训练代码中没有明确需求时，优先使用普通的 Out-of-place Operation：

x = x + 2
y = torch.relu(x)

不要仅仅为了节省显存，就把所有操作改成带下划线的版本。PyTorch 的 Autograd 会主动复用和释放缓冲区，原地操作通常带来的内存收益没有想象中明显。

11. 常用操作对照

需求	常用方法
根据数据创建	torch.tensor()
创建全 0、全 1 Tensor	torch.zeros()、torch.ones()
创建随机 Tensor	torch.rand()、torch.randn()、torch.randint()
创建相同形状 Tensor	torch.zeros_like()、torch.ones_like() 等
改变形状	view()、reshape()
交换维度	transpose()、permute()
增删大小为 1 的维度	squeeze()、unsqueeze()
扩展或重复	expand()、repeat_interleave()
拼接	torch.cat()
拆分	torch.split()、torch.chunk()
统计	torch.max()、torch.mean()、torch.numel()
选择前 \(k\) 个值	torch.topk()
创建坐标网格	torch.meshgrid()
按索引写入	torch.scatter()

参考资料

• PyTorch Tensor
• Tensor Views
• Tensor Attributes
• Autograd 中的 In-place Operation

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Aug 16, 2026 05:04 pm