Chuoer47's Blog

Appearance

张量操作详解 ​

数学运算 ​

逐元素运算 ​

逐元素运算是对张量中每个元素独立进行的运算，要求参与运算的张量形状相同（或可广播）。

python
import torch

a = torch.tensor([1.0, 2.0, 3.0, 4.0])
b = torch.tensor([2.0, 3.0, 4.0, 5.0])

# 加法
c = a + b             # 运算符写法
c = torch.add(a, b)   # 函数写法
print(f"加法: {c}")     # tensor([3., 5., 7., 9.])

# 减法
c = a - b
c = torch.sub(a, b)
print(f"减法: {c}")     # tensor([-1., -1., -1., -1.])

# 乘法（逐元素）
c = a * b
c = torch.mul(a, b)
print(f"乘法: {c}")     # tensor([ 2.,  6., 12., 20.])

# 除法
c = a / b
c = torch.div(a, b)
print(f"除法: {c}")     # tensor([0.5000, 0.6667, 0.7500, 0.8000])

# 整除
c = torch.div(b, a, rounding_mode="trunc")
print(f"整除: {c}")     # tensor([2., 1., 1., 1.])

# 取余
c = torch.remainder(b, a)
print(f"取余: {c}")     # tensor([0., 1., 1., 1.])

# 幂运算
c = a ** 2
c = torch.pow(a, 2)
print(f"幂运算: {c}")   # tensor([ 1.,  4.,  9., 16.])

# 平方根
c = torch.sqrt(a)
print(f"平方根: {c}")    # tensor([1.0000, 1.4142, 1.7321, 2.0000])

# 指数
c = torch.exp(a)
print(f"指数: {c}")      # tensor([ 2.7183,  7.3891, 20.0855, 54.5981])

# 对数
c = torch.log(a)
print(f"自然对数: {c}")   # tensor([0.0000, 0.6931, 1.0986, 1.3863])

# 绝对值
c = torch.abs(torch.tensor([-1, -2, 3]))
print(f"绝对值: {c}")    # tensor([1, 2, 3])

# 取反
c = -a
c = torch.neg(a)
print(f"取反: {c}")      # tensor([-1., -2., -3., -4.])

矩阵乘法 ​

PyTorch 提供了多种矩阵乘法方式，适用于不同的维度场景。

python
import torch

# ========== 2D 矩阵乘法 ==========
A = torch.tensor([[1.0, 2.0],
                   [3.0, 4.0]])  # 形状 (2, 2)

B = torch.tensor([[5.0, 6.0],
                   [7.0, 8.0]])  # 形状 (2, 2)

# 方法一: torch.matmul 或 @ 运算符
C = torch.matmul(A, B)
C = A @ B
print(f"矩阵乘法 (2x2 @ 2x2):\n{C}")
# tensor([[19., 22.],
#         [43., 50.]])

# 方法二: torch.mm（仅用于 2D 矩阵）
C = torch.mm(A, B)
print(f"torch.mm:\n{C}")

# 方法三: torch.bmm（用于 3D 批量矩阵乘法）
batch_A = torch.randn(3, 2, 4)  # 3 个 2x4 矩阵
batch_B = torch.randn(3, 4, 5)  # 3 个 4x5 矩阵
batch_C = torch.bmm(batch_A, batch_B)
print(f"批量矩阵乘法形状: {batch_C.shape}")  # torch.Size([3, 2, 5])

# ========== 高维矩阵乘法（广播） ==========
# matmul 支持广播，可以处理高维张量
x = torch.randn(2, 3, 4)  # 形状 (2, 3, 4)
y = torch.randn(2, 4, 5)  # 形状 (2, 4, 5)
z = torch.matmul(x, y)    # 形状 (2, 3, 5)
print(f"高维 matmul 形状: {z.shape}")

# 广播场景
x = torch.randn(2, 1, 3, 4)  # 形状 (2, 1, 3, 4)
y = torch.randn(   5, 4, 2)  # 形状 (5, 4, 2)
z = x @ y  # 广播后: (2, 5, 3, 2)
print(f"广播 matmul 形状: {z.shape}")

# ========== 向量运算 ==========
v1 = torch.tensor([1.0, 2.0, 3.0])
v2 = torch.tensor([4.0, 5.0, 6.0])

# 点积（内积）
dot = torch.dot(v1, v2)
print(f"点积: {dot}")  # tensor(32.)
# 等价于 (v1 * v2).sum()

# 外积
outer = torch.outer(v1, v2)
print(f"外积形状: {outer.shape}")  # torch.Size([3, 3])

广播机制详解 ​

广播（Broadcasting）是 PyTorch 中一个强大的特性，允许在不同形状的张量之间进行算术运算，无需显式地复制数据。

广播规则：

  1. 如果两个张量的维度数不同，较小维度的张量会在左侧补 1
  2. 如果两个张量在某个维度上的大小不同，且其中一个为 1，则该维度会扩展为另一个张量的大小
  3. 如果两个张量在某个维度上的大小不同且都不为 1，则无法广播，会报错
python
import torch

# ========== 广播示例 ==========

# 示例 1: 标量与张量
a = torch.tensor([1, 2, 3])
b = 2
print(f"张量 + 标量: {a + b}")  # tensor([3, 4, 5])
# 标量 2 被广播为 [2, 2, 2]

# 示例 2: 不同维度的张量
# 形状 (3, 4) 和 (4,)
x = torch.ones(3, 4)
y = torch.tensor([1, 2, 3, 4])
print(f"(3,4) + (4,) 结果形状: {(x + y).shape}")  # torch.Size([3, 4])
# y 的形状从 (4,) 广播为 (1, 4) 再扩展为 (3, 4)

# 示例 3: 维度补 1
# 形状 (3, 1) 和 (1, 4)
a = torch.tensor([[1], [2], [3]])  # 形状 (3, 1)
b = torch.tensor([[10, 20, 30, 40]])  # 形状 (1, 4)
c = a + b
print(f"(3,1) + (1,4) 结果形状: {c.shape}")  # torch.Size([3, 4])
print(f"结果:\n{c}")
# tensor([[11, 21, 31, 41],
#         [12, 22, 32, 42],
#         [13, 23, 33, 43]])

# 示例 4: 无法广播的情况
# 形状 (3, 4) 和 (3, 5) 无法广播
try:
    a = torch.ones(3, 4)
    b = torch.ones(3, 5)
    c = a + b
except RuntimeError as e:
    print(f"广播错误: {e}")

# 示例 5: 多维广播
# 形状 (2, 3, 1) 和 (1, 1, 4)
a = torch.ones(2, 3, 1)
b = torch.ones(1, 1, 4)
c = a + b
print(f"(2,3,1) + (1,1,4) = {c.shape}")  # torch.Size([2, 3, 4])

# ========== 手动扩展广播 ==========
# expand: 不复制数据的扩展（只读）
a = torch.tensor([[1], [2], [3]])  # 形状 (3, 1)
b = a.expand(3, 4)  # 扩展为 (3, 4)
print(f"expand 结果:\n{b}")

# repeat: 复制数据的扩展
c = a.repeat(1, 4)  # 在第 1 维重复 4 次
print(f"repeat 结果:\n{c}")

# 区别: expand 更高效（共享内存），repeat 会实际复制

比较运算 ​

比较运算返回布尔类型的张量，常用于条件判断、掩码生成和索引筛选。

python
import torch

a = torch.tensor([1, 2, 3, 4, 5])
b = torch.tensor([1, 3, 3, 2, 5])

# 等于
eq = a == b
eq = torch.eq(a, b)
print(f"等于: {eq}")      # tensor([ True, False,  True, False,  True])

# 不等于
ne = a != b
ne = torch.ne(a, b)
print(f"不等于: {ne}")     # tensor([False,  True, False,  True, False])

# 大于
gt = a > b
gt = torch.gt(a, b)
print(f"大于: {gt}")       # tensor([False, False, False,  True, False])

# 小于
lt = a < b
lt = torch.lt(a, b)
print(f"小于: {lt}")       # tensor([False,  True, False, False, False])

# 大于等于
ge = a >= b
ge = torch.ge(a, b)
print(f"大于等于: {ge}")   # tensor([ True, False,  True,  True,  True])

# 小于等于
le = a <= b
le = torch.le(a, b)
print(f"小于等于: {le}")   # tensor([ True,  True,  True, False,  True])

# 结合布尔运算
mask = (a > 2) & (a < 5)  # 逻辑与
print(f"2 < a < 5 的掩码: {mask}")  # tensor([False, False,  True,  True, False])

mask = (a < 2) | (a > 4)  # 逻辑或
print(f"a < 2 或 a > 4: {mask}")   # tensor([ True, False, False, False,  True])

mask = ~(a == 3)  # 逻辑非
print(f"a != 3: {mask}")           # tensor([ True,  True, False,  True,  True])

# any() 和 all()
x = torch.tensor([True, False, True])
print(f"any: {x.any()}")  # tensor(True)
print(f"all: {x.all()}")  # tensor(False)

# 排序后的索引比较
sorted_vals, sorted_idx = a.sort()
print(f"排序: {sorted_vals}, 索引: {sorted_idx}")

索引与切片 ​

基础索引与切片 ​

基础索引与 Python 列表和 NumPy 的切片语法一致，但返回的是**视图（view）**而非副本。

python
import torch

# 创建一个 2D 张量
x = torch.tensor([[1, 2, 3, 4],
                   [5, 6, 7, 8],
                   [9, 10, 11, 12]])
print(f"原始张量:\n{x}")

# 单元素索引
print(f"x[0, 0]: {x[0, 0]}")     # tensor(1)
print(f"x[1, 2]: {x[1, 2]}")     # tensor(7)

# 行切片
print(f"x[0]: {x[0]}")           # tensor([1, 2, 3, 4])
print(f"x[0:2]:\n{x[0:2]}")      # 前两行
print(f"x[-1]: {x[-1]}")         # tensor([9, 10, 11, 12]) 最后一行

# 列切片
print(f"x[:, 0]: {x[:, 0]}")     # tensor([1, 5, 9])
print(f"x[:, 1:3]:\n{x[:, 1:3]}") # 第 1、2 列

# 行列组合切片
print(f"x[0:2, 1:3]:\n{x[0:2, 1:3]}")

# 步长切片
print(f"x[::2]:\n{x[::2]}")      # 隔行取
print(f"x[:, ::2]:\n{x[:, ::2]}") # 隔列取

# 基础索引返回视图（共享内存）
sub = x[0:2, 0:2]
sub[0, 0] = 999
print(f"修改子张量后原始张量也变了:\n{x}")
# 注意: 这是视图操作，修改 sub 会影响 x

高级索引 ​

高级索引使用整数数组或张量进行索引，返回的是副本而非视图。

python
import torch

x = torch.tensor([[10, 20, 30],
                   [40, 50, 60],
                   [70, 80, 90]])

# 花式索引（Fancy Indexing）：使用整数列表/张量索引
# 选择特定行
rows = torch.tensor([0, 2])
print(f"选择第 0、2 行:\n{x[rows]}")
# tensor([[10, 20, 30],
#         [70, 80, 90]])

# 选择特定列
cols = torch.tensor([0, 2])
print(f"选择第 0、2 列:\n{x[:, cols]}")

# 行列组合索引
print(f"x[[0, 1, 2], [0, 1, 2]]: {x[[0, 1, 2], [0, 1, 2]]}")
# 等价于选择 x[0,0], x[1,1], x[2,2]
# tensor([10, 50, 90])

# 使用 torch.index_select
# 按维度选择
selected = torch.index_select(x, dim=0, index=torch.tensor([0, 2]))
print(f"index_select(dim=0):\n{selected}")

selected = torch.index_select(x, dim=1, index=torch.tensor([1]))
print(f"index_select(dim=1):\n{selected}")

# 使用 gather
# gather 在指定维度上按索引收集值
src = torch.tensor([[1, 2, 3],
                     [4, 5, 6]])
# dim=1 时，index 的形状决定输出形状
idx = torch.tensor([[0, 1],
                     [2, 0]])
result = torch.gather(src, dim=1, index=idx)
print(f"gather 结果:\n{result}")
# tensor([[1, 2],
#         [6, 4]])
# 第一行: 取 src[0,0]=1, src[0,1]=2
# 第二行: 取 src[1,2]=6, src[1,0]=4

# 使用 scatter（gather 的逆操作）
result = torch.zeros(2, 3, dtype=torch.long)
idx = torch.tensor([[0, 1, 2],
                     [2, 0, 1]])
src_vals = torch.tensor([[10, 20, 30],
                          [40, 50, 60]])
result.scatter_(dim=1, index=idx, src=src_vals)
print(f"scatter 结果:\n{result}")
# tensor([[10, 20, 30],
#         [50, 60, 40]])

布尔索引 ​

布尔索引（Boolean Indexing）使用布尔张量作为掩码来选择元素。

python
import torch

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

# 布尔掩码
mask = x > 4
print(f"大于 4 的掩码:\n{mask}")

# 使用布尔掩码选择元素（返回 1D 张量）
selected = x[mask]
print(f"大于 4 的元素: {selected}")  # tensor([5, 6, 7, 8, 9])

# 直接使用布尔表达式
print(f"偶数: {x[x % 2 == 0]}")    # tensor([2, 4, 6, 8])

# 布尔索引结合条件
print(f"大于 3 且小于 8 的元素: {x[(x > 3) & (x < 8)]}")

# 使用 torch.where（三元操作）
# torch.where(condition, x, y) 满足条件取 x，否则取 y
y = torch.where(x > 5, x, torch.zeros_like(x))
print(f"大于 5 保留，否则置 0:\n{y}")

# 用 where 实现 ReLU 激活函数
a = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
relu = torch.where(a > 0, a, torch.zeros_like(a))
print(f"ReLU: {relu}")  # tensor([0., 0., 0., 1., 2.])

# torch.masked_select
mask = x > 5
result = torch.masked_select(x, mask)
print(f"masked_select: {result}")  # tensor([6, 7, 8, 9])

# 使用 nonzero 获取非零元素的索引
nz = torch.tensor([0, 3, 0, 7, 0]).nonzero()
print(f"非零元素索引: {nz.squeeze()}")

形状操作 ​

reshape 和 view ​

python
import torch

x = torch.arange(12)
print(f"原始: {x}")  # tensor([ 0,  1,  2, ..., 11])

# reshape: 可以返回视图或副本（取决于内存是否连续）
y = x.reshape(3, 4)
print(f"reshape(3, 4):\n{y}")

# 使用 -1 自动推断维度
y = x.reshape(3, -1)  # -1 自动计算为 4
print(f"reshape(3, -1):\n{y}")

y = x.reshape(-1, 4)  # -1 自动计算为 3
print(f"reshape(-1, 4):\n{y}")

# view: 要求张量在内存中连续（更严格）
z = x.view(3, 4)
print(f"view(3, 4):\n{z}")

# view 的限制
t = torch.randn(2, 3)
t_t = t.t()  # 转置后内存不连续
try:
    t_t.view(6)
except RuntimeError as e:
    print(f"view 错误: {e}")

# 解决方法: 先 contiguous() 再 view
t_contiguous = t_t.contiguous()
print(f"contiguous 后 view: {t_contiguous.view(6)}")

# reshape 和 view 的区别:
# - view 要求连续内存，不连续时会报错
# - reshape 不要求连续内存，不连续时会返回副本
# - 两者在连续内存时都返回视图

transpose 和 permute ​

python
import torch

# transpose: 交换两个维度
x = torch.randn(2, 3, 4)
y = x.transpose(0, 1)  # 交换第 0 维和第 1 维
print(f"transpose(0,1) 形状: {y.shape}")  # torch.Size([3, 2, 4])

y = x.transpose(0, 2)  # 交换第 0 维和第 2 维
print(f"transpose(0,2) 形状: {y.shape}")  # torch.Size([4, 3, 2])

# 2D 转置
m = torch.tensor([[1, 2, 3],
                   [4, 5, 6]])
print(f"转置:\n{m.t()}")  # 等价于 m.transpose(0, 1)
print(f"转置:\n{m.T}")

# permute: 同时重新排列所有维度
x = torch.randn(2, 3, 4, 5)  # (N, C, H, W)
y = x.permute(0, 2, 3, 1)     # -> (N, H, W, C)
print(f"permute(0,2,3,1) 形状: {y.shape}")  # torch.Size([2, 4, 5, 3])

# 常见用法: 图像从 (C, H, W) 转为 (H, W, C)
img = torch.randn(3, 224, 224)  # (C, H, W)
img_hwc = img.permute(1, 2, 0)  # (H, W, C)
print(f"图像格式转换: {img.shape} -> {img_hwc.shape}")

squeeze 和 unsqueeze ​

python
import torch

# squeeze: 移除大小为 1 的维度
x = torch.randn(1, 3, 1, 4)
print(f"原始形状: {x.shape}")  # torch.Size([1, 3, 1, 4])

y = x.squeeze()  # 移除所有大小为 1 的维度
print(f"squeeze() 形状: {y.shape}")  # torch.Size([3, 4])

y = x.squeeze(0)  # 仅移除第 0 维
print(f"squeeze(0) 形状: {y.shape}")  # torch.Size([3, 1, 4])

y = x.squeeze(2)  # 仅移除第 2 维
print(f"squeeze(2) 形状: {y.shape}")  # torch.Size([1, 3, 4])

# unsqueeze: 在指定位置插入大小为 1 的维度
x = torch.tensor([1, 2, 3, 4])
print(f"原始形状: {x.shape}")  # torch.Size([4])

y = x.unsqueeze(0)  # 在第 0 维插入
print(f"unsqueeze(0) 形状: {y.shape}")  # torch.Size([1, 4])
print(f"unsqueeze(0):\n{y}")

y = x.unsqueeze(1)  # 在第 1 维插入
print(f"unsqueeze(1) 形状: {y.shape}")  # torch.Size([4, 1])
print(f"unsqueeze(1):\n{y}")

# 常见用法: 给数据添加 batch 维度
single_img = torch.randn(3, 224, 224)  # 单张图像
batch_img = single_img.unsqueeze(0)    # 添加 batch 维度
print(f"添加 batch 维度: {single_img.shape} -> {batch_img.shape}")
# torch.Size([3, 224, 224]) -> torch.Size([1, 3, 224, 224])

flatten 和 cat / stack ​

python
import torch

# flatten: 将张量展平
x = torch.randn(2, 3, 4)
y = x.flatten()  # 展平为 1D
print(f"flatten() 形状: {y.shape}")  # torch.Size([24])

y = x.flatten(start_dim=1)  # 从第 1 维开始展平
print(f"flatten(1) 形状: {y.shape}")  # torch.Size([2, 12])

y = x.flatten(start_dim=0, end_dim=1)  # 展平第 0-1 维
print(f"flatten(0,1) 形状: {y.shape}")  # torch.Size([6, 4])

# cat: 沿现有维度拼接
a = torch.randn(2, 3)
b = torch.randn(2, 3)

c = torch.cat([a, b], dim=0)  # 沿第 0 维拼接
print(f"cat dim=0 形状: {c.shape}")  # torch.Size([4, 3])

c = torch.cat([a, b], dim=1)  # 沿第 1 维拼接
print(f"cat dim=1 形状: {c.shape}")  # torch.Size([2, 6])

# 多个张量拼接
c = torch.cat([a, b, a], dim=0)
print(f"三个张量 cat: {c.shape}")  # torch.Size([6, 3])

# stack: 沿新维度堆叠
c = torch.stack([a, b], dim=0)  # 在第 0 维堆叠
print(f"stack dim=0 形状: {c.shape}")  # torch.Size([2, 2, 3])

c = torch.stack([a, b], dim=1)  # 在第 1 维堆叠
print(f"stack dim=1 形状: {c.shape}")  # torch.Size([2, 2, 3])

c = torch.stack([a, b], dim=2)  # 在第 2 维堆叠
print(f"stack dim=2 形状: {c.shape}")  # torch.Size([2, 3, 2])

# cat vs stack 区别:
# cat: 沿已有维度拼接，不增加维度数
# stack: 沿新维度堆叠，增加一个维度

split 和 chunk ​

python
import torch

x = torch.arange(10).float()
print(f"原始: {x}")

# split: 按指定大小拆分
parts = x.split(3)  # 每份 3 个元素（最后一份可以更少）
for i, p in enumerate(parts):
    print(f"split(3) 第 {i} 部分: {p}")
# tensor([0., 1., 2.])
# tensor([3., 4., 5.])
# tensor([6., 7., 8.])
# tensor([9.])

# 指定各部分大小
parts = x.split([2, 3, 5])
for i, p in enumerate(parts):
    print(f"split([2,3,5]) 第 {i} 部分: {p}")

# chunk: 按数量拆分
parts = x.chunk(4)  # 拆成 4 份
for i, p in enumerate(parts):
    print(f"chunk(4) 第 {i} 部分: {p}")

# 2D 张量的 split
y = torch.randn(6, 4)
parts = y.split(2, dim=0)  # 沿第 0 维拆分，每份 2 行
for i, p in enumerate(parts):
    print(f"split(2, dim=0) 第 {i} 部分形状: {p.shape}")

repeat 和 expand ​

python
import torch

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

# repeat: 沿各维度重复（数据被实际复制）
y = x.repeat(3, 1)  # 第 0 维重复 3 次，第 1 维重复 1 次
print(f"repeat(3, 1):\n{y}")
# tensor([[1, 2, 3],
#         [1, 2, 3],
#         [1, 2, 3]])

y = x.repeat(3, 2)  # 第 0 维重复 3 次，第 1 维重复 2 次
print(f"repeat(3, 2):\n{y}")

# expand: 扩展大小为 1 的维度（不复制数据，更高效）
x = torch.tensor([[1, 2, 3]])  # 形状 (1, 3)
y = x.expand(3, 3)  # 扩展第 0 维
print(f"expand(3, 3):\n{y}")
# tensor([[1, 2, 3],
#         [1, 2, 3],
#         [1, 2, 3]])

# expand 的限制: 只能扩展大小为 1 的维度
x = torch.tensor([[1, 2, 3]])
try:
    y = x.expand(3, 4)  # 第 1 维大小为 3，不能扩展到 4
except RuntimeError as e:
    print(f"expand 错误: {e}")

# 使用 -1 保持该维度不变
x = torch.randn(1, 3, 1, 5)
y = x.expand(2, -1, 4, -1)  # -1 表示保持原大小
print(f"expand 后形状: {y.shape}")  # torch.Size([2, 3, 4, 5])

# repeat 与 expand 的区别:
# - expand 不复制数据，返回视图（只读），更高效
# - repeat 实际复制数据，占用更多内存
# - expand 只能扩展大小为 1 的维度，repeat 可以任意重复

归约操作 ​

sum、mean、max、min ​

python
import torch

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

# sum: 求和
print(f"总和: {x.sum()}")             # tensor(21.)
print(f"按列求和: {x.sum(dim=0)}")     # tensor([5., 7., 9.])
print(f"按行求和: {x.sum(dim=1)}")     # tensor([ 6., 15.])

# 保持维度
print(f"按列求和（保持维度）:\n{x.sum(dim=0, keepdim=True)}")
# tensor([[5., 7., 9.]])  形状 (1, 3)

# mean: 均值
print(f"均值: {x.mean()}")            # tensor(3.5000)
print(f"按列均值: {x.mean(dim=0)}")    # tensor([2.5000, 3.5000, 4.5000])

# max / min: 最大/最小值
print(f"最大值: {x.max()}")            # tensor(6.)
print(f"每列最大值: {x.max(dim=0)}")   # 返回 (values, indices)

# 分别获取值和索引
max_vals, max_indices = x.max(dim=1)
print(f"每行最大值: {max_vals}")        # tensor([3., 6.])
print(f"每行最大值索引: {max_indices}") # tensor([2, 2])

# min 同理
min_vals, min_indices = x.min(dim=0)
print(f"每列最小值: {min_vals}")        # tensor([1., 2., 3.])

# prod: 连乘
print(f"总乘积: {x.prod()}")           # tensor(720.)

argmax、argmin、topk、sort ​

python
import torch

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

# argmax: 最大值的索引
print(f"全局 argmax: {x.argmax()}")        # tensor(5) 展平后的索引
print(f"每行 argmax: {x.argmax(dim=1)}")   # tensor([2, 2, 1])
print(f"每列 argmax: {x.argmax(dim=0)}")   # tensor([0, 2, 1])

# argmin: 最小值的索引
print(f"每行 argmin: {x.argmin(dim=1)}")   # tensor([1, 0, 0])

# topk: 获取最大的 k 个值
values, indices = x.topk(2, dim=1)
print(f"每行最大的 2 个值:\n{values}")
print(f"对应索引:\n{indices}")

# 最小的 k 个值
values, indices = x.topk(2, dim=1, largest=False)
print(f"每行最小的 2 个值:\n{values}")

# sort: 排序
sorted_vals, sorted_indices = x.sort(dim=1)
print(f"每行排序:\n{sorted_vals}")
print(f"排序索引:\n{sorted_indices}")

# 降序排序
sorted_vals, sorted_indices = x.sort(dim=1, descending=True)
print(f"每行降序排序:\n{sorted_vals}")

# argsort: 仅返回排序索引
indices = x.argsort(dim=1)
print(f"argsort 结果:\n{indices}")

原地操作与非原地操作 ​

原地操作概述 ​

在 PyTorch 中，原地操作（In-place Operation）通过在方法名后加下划线 _ 来标识。原地操作直接修改张量本身，而非返回新的张量。

python
import torch

# ========== 非原地操作（返回新张量） ==========
x = torch.tensor([1.0, 2.0, 3.0])
y = x.add(1)  # 返回新张量
print(f"x 不变: {x}")       # tensor([1., 2., 3.])
print(f"新张量 y: {y}")     # tensor([2., 3., 4.])

# ========== 原地操作（修改原始张量） ==========
x = torch.tensor([1.0, 2.0, 3.0])
x.add_(1)  # 原地加 1
print(f"原地操作后 x: {x}")  # tensor([2., 3., 4.])

# 常见的原地操作
x = torch.tensor([1.0, 2.0, 3.0, 4.0])

x.add_(10)          # 原地加法
print(f"add_: {x}")

x.mul_(2)           # 原地乘法
print(f"mul_: {x}")

x.sub_(5)           # 原地减法
print(f"sub_: {x}")

x.div_(3)           # 原地除法
print(f"div_: {x}")

x.clamp_(min=0)     # 原地裁剪（ReLU）
print(f"clamp_: {x}")

x.zero_()           # 原地置零
print(f"zero_: {x}")

x.fill_(3.14)       # 原地填充
print(f"fill_: {x}")

x.abs_()            # 原地取绝对值
x.relu_()           # 原地 ReLU
x.sigmoid_()        # 原地 sigmoid
x.tanh_()           # 原地 tanh

原地操作的注意事项 ​

python
import torch

# ========== 注意 1: requires_grad=True 的张量不能原地操作 ==========
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
try:
    x.add_(1)
except RuntimeError as e:
    print(f"错误: {e}")
# RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

# ========== 注意 2: 原地操作可能破坏计算图 ==========
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * 2
# z = y + x  # 这样是安全的
# x.add_(1)  # 如果原地修改 x，y 的计算历史会被破坏

# ========== 注意 3: 切片原地赋值是允许的 ==========
x = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
y = x * 2
y[0] = 100  # 切片赋值是允许的
# 或
y.detach()[0] = 100

# ========== 最佳实践 ==========
# 1. 在不需要梯度追踪时使用原地操作节省内存
# 2. 在模型训练循环中，避免对需要梯度的张量使用原地操作
# 3. 对于叶子张量（leaf tensor），使用 .data 属性进行原地操作
x = torch.tensor([1.0, 2.0], requires_grad=True)
x.data.add_(1)  # 允许，但会跳过梯度追踪
print(f"data.add_ 后: {x}")

实用技巧 ​

einsum: 爱因斯坦求和 ​

torch.einsum 是一个极其灵活的函数，可以用简洁的字符串表达各种张量运算。

python
import torch

A = torch.randn(3, 4)
B = torch.randn(4, 5)

# 矩阵乘法: C_ij = A_ik * B_kj
C = torch.einsum('ik,kj->ij', A, B)
print(f"矩阵乘法形状: {C.shape}")  # torch.Size([3, 5])

# 向量点积
v1 = torch.randn(5)
v2 = torch.randn(5)
dot = torch.einsum('i,i->', v1, v2)
print(f"点积: {dot}")

# 矩阵转置
T = torch.einsum('ij->ji', A)
print(f"转置形状: {T.shape}")  # torch.Size([4, 3])

# 矩阵迹（对角线之和）
M = torch.randn(4, 4)
trace = torch.einsum('ii->', M)
print(f"矩阵的迹: {trace}")

# 批量矩阵乘法
batch_A = torch.randn(8, 3, 4)
batch_B = torch.randn(8, 4, 5)
batch_C = torch.einsum('bik,bkj->bij', batch_A, batch_B)
print(f"批量矩阵乘法形状: {batch_C.shape}")  # torch.Size([8, 3, 5])

# 注意力机制中的 QKV 计算
# Q: (batch, seq_len, d_k), K: (batch, seq_len, d_k), V: (batch, seq_len, d_v)
batch, seq_len, d_k, d_v = 2, 10, 64, 32
Q = torch.randn(batch, seq_len, d_k)
K = torch.randn(batch, seq_len, d_k)
V = torch.randn(batch, seq_len, d_v)

# 注意力分数: (batch, seq_len, seq_len)
attn = torch.einsum('bqd,bkd->bqk', Q, K) / (d_k ** 0.5)
attn = torch.softmax(attn, dim=-1)

# 加权求和: (batch, seq_len, d_v)
output = torch.einsum('bqk,bvd->bqd', attn, V)
print(f"注意力输出形状: {output.shape}")  # torch.Size([2, 10, 32])

# 张量缩并（trace of product）
A = torch.randn(3, 4)
B = torch.randn(3, 4)
result = torch.einsum('ij,ij->', A, B)  # Frobenius 内积
print(f"Frobenius 内积: {result}")

# 求对角线元素
M = torch.randn(4, 4)
diag = torch.einsum('ii->i', M)
print(f"对角线元素: {diag}")

where 和 clamp ​

python
import torch

# ========== where: 条件选择 ==========
x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])

# 基本用法: where(condition, x, y)
# condition 为 True 取 x，否则取 y
result = torch.where(x > 0, x, torch.tensor(0.0))
print(f"ReLU: {result}")  # tensor([0., 0., 0., 1., 2.])

# 两个张量之间选择
a = torch.tensor([1, 2, 3, 4, 5])
b = torch.tensor([10, 20, 30, 40, 50])
mask = torch.tensor([True, False, True, False, True])
result = torch.where(mask, a, b)
print(f"条件选择: {result}")  # tensor([10, 2, 30, 4, 50])

# 2D 张量示例
x = torch.randn(3, 4)
result = torch.where(x > 0, x, torch.tensor(0.0))
print(f"负值置零:\n{result}")

# ========== clamp: 裁剪 ==========
x = torch.tensor([0.5, 1.5, 2.5, 3.5, 4.5])

# 限制在 [min, max] 范围内
y = torch.clamp(x, min=1.0, max=3.0)
print(f"clamp(1.0, 3.0): {y}")  # tensor([1.0000, 1.5000, 2.5000, 3.0000, 3.0000])

# 仅设置最小值
y = torch.clamp(x, min=2.0)
print(f"clamp(min=2.0): {y}")   # tensor([2.0000, 2.0000, 2.5000, 3.5000, 4.5000])

# 仅设置最大值
y = torch.clamp(x, max=3.0)
print(f"clamp(max=3.0): {y}")   # tensor([0.5000, 1.5000, 2.5000, 3.0000, 3.0000])

# clamp_min 和 clamp_max 的简写
y = torch.clamp_min(x, 2.0)
y = torch.clamp_max(x, 3.0)

# 梯度裁剪中的应用
x = torch.randn(3, requires_grad=True)
y = x ** 2
y.sum().backward()
# 梯度裁剪
torch.nn.utils.clip_grad_norm_([x], max_norm=1.0)

unique: 唯一值操作 ​

python
import torch

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

# 基本用法: 返回唯一值
unique_vals = torch.unique(x)
print(f"唯一值: {unique_vals}")  # tensor([1, 2, 3, 4, 5])

# 排序后的唯一值（默认已排序）
unique_vals = torch.unique(x, sorted=True)
print(f"排序后的唯一值: {unique_vals}")

# 返回逆索引（用于重建原始张量）
unique_vals, inverse_indices = torch.unique(x, return_inverse=True)
print(f"唯一值: {unique_vals}")
print(f"逆索引: {inverse_indices}")
# 重建: unique_vals[inverse_indices] == x

# 返回每个唯一值的计数
unique_vals, counts = torch.unique(x, return_counts=True)
print(f"唯一值: {unique_vals}")
print(f"计数: {counts}")

# 2D 张量的 unique
x = torch.tensor([[1, 2], [3, 2], [1, 2]])
unique_rows = torch.unique(x, dim=0)
print(f"唯一行:\n{unique_rows}")

# 常见用途: 统计类别数
labels = torch.tensor([0, 1, 2, 1, 0, 2, 3, 1])
num_classes = torch.unique(labels).numel()
print(f"类别数: {num_classes}")  # 4

# 配合 where 使用: 获取某个类别的所有索引
target_class = 1
indices = torch.where(labels == target_class)[0]
print(f"类别 {target_class} 的索引: {indices}")

# 使用 bincount 统计各类别数量
counts = torch.bincount(labels)
print(f"各类别数量: {counts}")  # tensor([2, 3, 2, 1])

更多实用技巧 ​

python
import torch

# ========== masked_fill: 掩码填充 ==========
x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
mask = torch.tensor([True, False, True, False, True])
result = x.masked_fill(mask, value=0.0)
print(f"masked_fill: {result}")  # tensor([0., 2., 0., 4., 0.])

# 常见用途: 注意力掩码
seq_len = 5
# 创建上三角掩码（因果掩码）
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
print(f"因果掩码:\n{~mask}")

# ========== norm: 范数计算 ==========
x = torch.tensor([3.0, 4.0])
print(f"L1 范数: {x.norm(p=1)}")      # tensor(7.)
print(f"L2 范数: {x.norm(p=2)}")      # tensor(5.)
print(f"无穷范数: {x.norm(p=float('inf'))}")  # tensor(4.)

# 矩阵范数
M = torch.randn(3, 4)
print(f"Frobenius 范数: {M.norm()}")

# ========== cumsum / cumprod: 累积运算 ==========
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
print(f"累积和: {torch.cumsum(x, dim=0)}")   # tensor([ 1.,  3.,  6., 10.])
print(f"累积积: {torch.cumprod(x, dim=0)}")  # tensor([ 1.,  2.,  6., 24.])

# ========== cross_entropy 简化示例 ==========
# 在实际训练中常用的组合操作
logits = torch.randn(3, 5)  # 3 个样本，5 个类别
targets = torch.tensor([1, 3, 4])
loss = torch.nn.functional.cross_entropy(logits, targets)
print(f"交叉熵损失: {loss}")

# ========== detach: 从计算图中分离 ==========
x = torch.randn(3, requires_grad=True)
y = x * 2
z = y.detach()  # z 不再追踪梯度
print(f"y requires_grad: {y.requires_grad}")  # True
print(f"z requires_grad: {z.requires_grad}")  # False

参考资料 ​