PyTorch张量与自动求导

PyTorch 的两个核心能力：张量（tensor），可以放在 GPU 上运算的多维数组；自动求导（autograd），自动计算梯度。有了这两样，就可以用梯度下降训练任意的神经网络。

以下代码都用 import torch。

创建张量

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(x.shape, x.dtype, x.device)   # torch.Size([2, 2]) torch.float32 cpu
 
torch.zeros(2, 3)          # 全 0
torch.ones(2, 3)           # 全 1
torch.arange(6)            # tensor([0, 1, 2, 3, 4, 5])
torch.randn(2, 3)          # 标准正态分布
torch.full((2, 2), 7.0)    # 填充指定值
torch.zeros_like(x)        # 形状、类型、设备都与 x 相同

对支持的数据类型，torch.from_numpy(a) 与 NumPy 数组共享 CPU 内存；普通 t.numpy() 也可共享存储，但要求张量位于 CPU、不需要梯度，且类型和布局受 NumPy 支持。GPU 张量通常要先 t.detach().cpu().numpy()，其中搬到 CPU 的步骤会复制数据。

import numpy as np
a = np.arange(3)
t = torch.from_numpy(a)
a[0] = 100
print(t)                   # tensor([100,   1,   2])

形状操作

这是读 Transformer 代码时出现最多的一类操作。

x = torch.arange(24)
y = x.view(2, 3, 4)        # 改变形状，要求新形状与原来的尺寸、步幅兼容
x.reshape(4, -1)           # 同上，必要时会自动复制，更省心
y.transpose(1, 2).shape    # 交换两个维度：(2, 4, 3)
y.permute(2, 0, 1).shape   # 任意重排维度：(4, 2, 3)
 
v = torch.tensor([1.0, 2.0, 3.0])
v.unsqueeze(0).shape       # 插入一个长度为 1 的维度：(1, 3)
v.unsqueeze(1).shape       # (3, 1)
torch.zeros(1, 3, 1).squeeze().shape   # 去掉所有长度为 1 的维度：(3,)
 
torch.cat([torch.zeros(2, 3), torch.ones(2, 3)], dim=0).shape    # 沿已有维度拼接：(4, 3)
torch.stack([torch.zeros(2, 3), torch.ones(2, 3)], dim=0).shape  # 新建一个维度堆叠：(2, 2, 3)

view 和 reshape 的区别：view 共享原存储，要求新形状与原尺寸、步幅兼容，并非所有非连续张量都不能 view。transpose、permute 通常只改变步幅，可能使特定的 view 操作失效；这时可先 .contiguous()，或用必要时自动复制的 reshape。

t = y.transpose(1, 2)
t.is_contiguous()          # False
# t.view(-1)               # 取消注释会触发 RuntimeError：此例不能直接展平
t.contiguous().view(-1)    # 可以
t.reshape(-1)              # 也可以

一个典型例子是多头注意力里的「拆头」和「合头」：

B, T, H, D = 2, 5, 4, 8                     # 批次、序列长度、头数、每头维度
h = torch.randn(B, T, H * D)
heads = h.view(B, T, H, D).transpose(1, 2)  # (B, H, T, D)：每个头独立计算注意力
back = heads.transpose(1, 2).reshape(B, T, H * D)
torch.equal(back, h)                        # True

细节见注意力机制。

dtype 与 device

dtype：浮点张量默认 float32，整数张量默认 int64。训练大模型时常用低精度来省显存、提速：

dtype	特点
torch.float32	默认精度
torch.float16	最大有限值为 65504；既可能上溢，也可能因小梯度下溢而丢失信息；损失缩放主要缓解后者
torch.bfloat16	与 float32 同为 8 位指数，动态范围接近（最大有限值约 $3.39\times 10^{38}$），但有效精度更低；通常不需要损失缩放

f = torch.tensor([1.0, 2.0])
f.to(torch.bfloat16).dtype   # torch.bfloat16
f.half().dtype               # torch.float16
 
# bf16 只有 8 位有效二进制位（约 2 到 3 位十进制），1 + 0.001 会被舍掉
torch.tensor(1.0, dtype=torch.bfloat16) + torch.tensor(0.001, dtype=torch.bfloat16)
# tensor(1., dtype=torch.bfloat16)

混合精度训练的细节见优化器、学习率调度与混合精度。

device：张量在哪个设备上，运算就在哪里进行。参与同一次运算的张量必须在同一个设备上，否则报错。

device = "cuda" if torch.cuda.is_available() else "cpu"
z = torch.randn(3, 3, device=device)   # 直接在目标设备上创建
w = torch.randn(3, 3).to(device)       # 或者创建后搬过去
z @ w

自动求导

把张量的 requires_grad 设为 True，PyTorch 就会记录对它做的所有运算，构成一张计算图。对最终的标量结果调用 .backward()，就会沿着计算图反向传播，把梯度写进每个叶子张量的 .grad 属性。

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x
y.backward()
print(x.grad)             # tensor(7.)，即 dy/dx = 2x + 3 在 x=2 处的值

模型参数就是这样的叶子张量：

w = torch.randn(3, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
inp = torch.tensor([1.0, 2.0, 3.0])
out = (w * inp).sum() + b
out.backward()
print(w.grad)             # 等于 inp
print(b.grad)             # tensor([1.])

原理见神经网络与反向传播和微积分与梯度。

几个要注意的行为：

1. 无参数的 backward() 要求输出只有一个元素。对多元素输出也可以求导，但必须显式传入同形状的上游梯度，计算向量–雅可比积；把输出求和或求平均只是常见用法。

v = torch.tensor([1.0, 2.0], requires_grad=True)
# (v * 2).backward()      # 取消注释会报错：缺少上游梯度
(v * 2).backward(torch.ones_like(v))
print(v.grad)             # tensor([2., 2.])，等价于对输出求和后反传

2. 梯度会累加，不会自动清零。 每次 backward() 都把新梯度加到 .grad 上，所以训练循环里每一步都要先清零。

x = torch.tensor(2.0, requires_grad=True)
(x ** 2).backward(); print(x.grad)   # tensor(4.)
(x ** 2).backward(); print(x.grad)   # tensor(8.)：累加了
x.grad.zero_()                       # 手动清零；训练时用 optimizer.zero_grad()

这个特性也有用处：梯度累积就是故意连续做几次 backward() 再更新参数，用小批次模拟大批次，见PyTorch模型训练。

3. 反向所需的已保存中间量默认会被释放。若再次沿同一计算图反传还需要这些中间量，通常会报错；确需复用时可指定 retain_graph=True，但会增加内存占用。这不等于 Python 中所有图节点对象都立即消失，也不是每种简单运算都必然在第二次反传时报错。

关闭梯度记录

推理和评估时不需要梯度，关掉记录可以省内存、提速：

with torch.no_grad():
    y = v * 2
y.requires_grad            # False
 
d = (v * 2).detach()       # 从计算图中分离出一个不需要梯度的张量
d.requires_grad            # False
 
@torch.inference_mode()    # 比 no_grad 更彻底，推理时推荐
def predict(t):
    return t * 2

在 AI 中的作用

• 所有的模型参数、输入数据、中间激活都是张量。
• 自动求导让研究者只需要写前向计算，梯度由框架自动算出，这是深度学习框架最核心的价值。
• dtype 和 device 决定了显存占用和速度，是大模型训练与推理优化的起点。

参考资料

• PyTorch：Tensor.backward
• PyTorch：Tensor.view
• PyTorch：自动混合精度示例

相关

• Python与NumPy：形状、广播、索引的规则与 PyTorch 基本一致。
• PyTorch模型训练：把张量和自动求导组织成完整的训练流程。
• 神经网络与反向传播：自动求导背后的算法。
• 梯度下降：拿到梯度之后如何更新参数。