PyTorch学习笔记(6) : torch.autograd

PyTorch 自动微分 (torch.autograd) 学习笔记

参考文档：PyTorch 2.11 官方文档

原文链接：Automatic differentiation package - torch.autograd

一、概述

torch.autograd 是 PyTorch 的自动微分引擎，用于实现任意标量值函数的自动微分。

核心特点：

对现有代码改动极小，只需将 requires_grad=True 设置到需要计算梯度的张量
仅支持浮点类型张量：half, float, double, bfloat16
支持复数类型张量：cfloat, cdouble

二、核心函数

2.1 反向传播基础

函数	说明
`torch.autograd.backward()`	计算给定张量相对于图叶节点的梯度之和
`torch.autograd.grad()`	计算并返回输出相对于输入的梯度之和

三、前向模式自动微分（Forward-mode AD）

⚠️ Beta 版本：API 可能变化，功能覆盖还在完善中

函数/类	说明
`forward_ad.dual_level`	前向 AD 的上下文管理器，所有前向 AD 计算必须在此上下文中进行
`forward_ad.make_dual()`	将张量值与其切线关联，创建用于前向 AD 梯度计算的"对偶张量"
`forward_ad.unpack_dual()`	解包对偶张量，获取张量值和前向 AD 梯度
`forward_ad.enter_dual_level()`	进入新的前向梯度层级
`forward_ad.exit_dual_level()`	退出前向梯度层级
`UnpackedDualTensor`	`unpack_dual()` 返回的命名元组，包含对偶张量的原值和切线分量

四、函数式高级 API

⚠️ Beta 版本：性能优化还在进行中

用于计算 Jacobian、Hessian 等高级微分操作。

函数	说明
`functional.jacobian()`	计算给定函数的雅可比矩阵
`functional.hessian()`	计算给定标量函数的 Hessian 矩阵
`functional.vjp()`	计算向量 v 与函数雅可比矩阵的点积（Vector-Jacobian Product）
`functional.jvp()`	计算函数雅可比矩阵与向量 v 的点积（Jacobian-Vector Product）
`functional.vhp()`	计算向量 v 与标量函数 Hessian 矩阵的点积
`functional.hvp()`	计算标量函数 Hessian 矩阵与向量 v 的点积

使用技巧： 如果函数接受非张量参数或 requires_grad=False 的张量，可以用 lambda 捕获：

python 复制代码

functional.jacobian(lambda x: f(x, constant, flag=flag), input)

五、梯度布局管理

5.1 默认梯度布局

当非稀疏参数在 backward() 中接收到非稀疏梯度时：

场景	行为
`param.grad` 初始为 None + 内存非重叠且密集	创建与 `param` 步幅匹配的 `.grad`
`param.grad` 初始为 None + 其他情况	创建行优先连续步幅的 `.grad`
已有 `.grad` + `create_graph=False`	原地累加，保留原有步幅
已有 `.grad` + `create_graph=True`	替换为新张量，尝试匹配原有步幅

性能优化建议： 每次迭代前将 .grad 设为 None，而非调用 zero_grad()：

python 复制代码

for param in model.parameters():
    param.grad = None
loss.backward()

5.2 手动控制梯度布局

如需手动控制 .grad 的步幅：

在第一次 backward() 前赋值 param.grad = 指定步幅的零张量
永远不要将其设为 None
create_graph=False 时布局保证保留

六、原地操作（In-place Operations）

⚠️ 强烈建议避免使用！

Autograd 的激进缓冲区释放和重用已经非常高效
原地操作很少能显著降低内存使用
除非内存压力极大，否则不需要使用

原地操作正确性检查

张量会跟踪应用于它们的原地操作。如果检测到张量被保存用于反向传播但随后被原地修改，启动反向传播时会报错。

七、张量自动微分属性与方法

属性/方法	说明
`torch.Tensor.grad`	默认为 None，第一次 `backward()` 后变为张量
`torch.Tensor.requires_grad`	是否需要计算梯度
`torch.Tensor.is_leaf`	是否为叶节点（`requires_grad=False` 的张量默认为叶节点）
`torch.Tensor.backward()`	计算当前张量相对于图叶节点的梯度
`torch.Tensor.detach()`	返回从当前图分离的新张量
`torch.Tensor.detach_()`	将张量从创建它的图中分离，使其成为叶节点
`torch.Tensor.register_hook()`	注册反向传播钩子
`torch.Tensor.register_post_accumulate_grad_hook()`	注册梯度累加后的钩子
`torch.Tensor.retain_grad()`	启用此张量的 `grad` 在 `backward()` 期间填充

八、自定义自动微分函数

8.1 Function 基类

python 复制代码

class torch.autograd.Function(*args, **kwargs)

创建自定义函数的步骤：

继承 Function 类
实现 forward() 和 backward() 静态方法
通过 apply 类方法调用（不要直接调用 forward()）

示例代码：

python 复制代码

class Exp(Function):
    @staticmethod
    def forward(ctx, i):
        result = i.exp()
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_tensors
        return grad_output * result

# 使用 apply 方法调用
output = Exp.apply(input)

8.2 Function 核心方法

方法	说明
`Function.forward()`	定义自定义函数的前向传播
`Function.backward()`	用反向模式自动微分定义操作的微分公式
`Function.jvp()`	用前向模式自动微分定义操作的微分公式
`Function.vmap()`	定义此函数在 `torch.vmap()` 下的行为

8.3 上下文方法（ctx）

方法	说明
`ctx.mark_dirty()`	标记给定张量为原地操作修改
`ctx.mark_non_differentiable()`	标记输出为不可微分
`ctx.save_for_backward()`	保存给定张量供后续 `backward()` 调用
`ctx.set_materialize_grads()`	设置是否实例化梯度张量

九、数值梯度检查

函数	说明
`gradcheck()`	用小有限差分检查梯度，对比解析梯度
`gradgradcheck()`	检查梯度的梯度
`GradcheckError`	`gradcheck()` 和 `gradgradcheck()` 抛出的错误

十、性能分析器（Profiler）

10.1 profile 类

python 复制代码

class torch.autograd.profiler.profile(
    enabled=True, *, 
    use_cuda=False,           # 将弃用
    use_device=None,          # 可选: 'cuda', 'xpu', 'mtia', 'privateuseone'
    record_shapes=False,      # 记录输入维度信息
    with_flops=False,         # 估计 FLOPs（仅支持矩阵乘法和2D卷积）
    profile_memory=False,     # 跟踪张量内存分配/释放
    with_stack=False,         # 记录源信息（文件和行号）
    with_modules=False,         # 记录模块层次结构
    use_kineto=False,         # 实验性：启用 Kineto 分析器
    use_cpu=True,             # 分析 CPU 事件
    experimental_config=None,
    acc_events=False,         # 跨多个分析周期累加 FunctionEvents
    post_processing_timeout_s=None  # 后处理超时时间
)

⚠️ 重要警告：

启用内存分析或源属性会增加额外开销
不要递归调用（不允许嵌套实例）
CUDA 多进程限制：use_device='cuda' 不能与 num_workers > 0 的 DataLoader 一起使用

使用示例：

python 复制代码

x = torch.randn((1, 1), requires_grad=True)
with torch.autograd.profiler.profile() as prof:
    for _ in range(100):
        y = x ** 2
        y.backward()

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

10.2 其他分析器

类/函数	说明
`emit_nvtx()`	使每个 autograd 操作发出 NVTX 范围，用于 nvprof
`emit_itt()`	使每个 autograd 操作发出 ITT 范围，用于 Intel VTune Profiler
`load_nvprof()`	打开并解析 nvprof 跟踪文件

十一、调试与异常检测

11.1 detect_anomaly

python 复制代码

class torch.autograd.detect_anomaly(check_nan=True)

功能：

启用前向传播检测，允许反向传播打印创建失败反向函数的前向操作 traceback
check_nan=True 时，任何产生 "nan" 的反向计算都会报错

⚠️ 仅用于调试，会显著减慢程序执行

11.2 set_detect_anomaly

python 复制代码

class torch.autograd.set_detect_anomaly(mode, check_nan=True)

可作为上下文管理器或函数使用，启用/禁用异常检测。

十二、计算图操作

12.1 图节点（Node）

grad_fn 属性持有 torch.autograd.graph.Node 对象（如果张量是 autograd 记录操作的输出）。

属性/方法	说明
`Node.name`	返回节点名称
`Node.metadata`	返回元数据
`Node.next_functions`	下一个函数
`Node.register_hook()`	注册反向传播钩子
`Node.register_prehook()`	注册反向传播前钩子

12.2 保存张量钩子

python 复制代码

class torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

用途： 定义中间结果的打包/解包方式，常用于用计算换内存（如保存到磁盘或 CPU 而非 GPU）。

钩子签名：

python 复制代码

pack_hook(tensor: Tensor) -> Any
unpack_hook(Any) -> Tensor

示例：

python 复制代码

def pack_hook(x):
    return x.cpu()  # 保存到 CPU

def unpack_hook(x):
    return x.to("cuda")  # 移回 GPU

with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    y = model(x)

12.3 其他图工具

类/函数	说明
`save_on_cpu(pin_memory=False, device_type='cuda')`	前向保存的张量存到 CPU，反向时取回
`disable_saved_tensors_hooks(error_message)`	禁用保存张量默认钩子功能
`register_multi_grad_hook(tensors, fn, mode='all')`	注册多梯度反向钩子（"all" 或 "any" 模式）
`allow_mutation_on_saved_tensors()`	允许修改保存用于反向传播的张量
`GradientEdge`	表示计算图中给定梯度边的对象
`get_gradient_edge(tensor)`	获取计算给定张量梯度的梯度边
`set_warn_on_accumulate_grad_stream_mismatch(enabled)`	设置当 AccumulateGrad 节点流与产生输入梯度的节点流不匹配时是否警告

十三、已弃用内容

Variable（已弃用）

Variable API 已完全弃用，现在：

Variable(tensor) 和 Variable(tensor, requires_grad) 返回张量而非 Variable
var.data 等同于 tensor.data
var.backward(), var.detach(), var.register_hook() 等方法直接用于张量

现代用法：

python 复制代码

autograd_tensor = torch.randn((2, 3, 4), requires_grad=True)

十四、学习要点总结

要点	说明
1. 自动微分核心	`requires_grad=True` + `backward()`
2. 避免原地操作	除非内存极度紧张，否则不建议使用
3. 梯度清零优化	用 `param.grad = None` 替代 `zero_grad()` 可能提升性能
4. 自定义函数	继承 `Function`，实现 `forward`/`backward`，用 `apply` 调用
5. 调试工具	`detect_anomaly` 用于定位梯度问题
6. 内存优化	`save_on_cpu` 和 `saved_tensors_hooks` 用于大模型训练
7. 性能分析	`profile` 类用于定位性能瓶颈

📌 建议： 实际使用时结合 PyTorch 官方教程进行代码实践，加深理解。