PyTorch自动微分模块：从原理到实战 ①
-
- [📋 目录](#📋 目录)
- [📝 摘要](#📝 摘要)
- [第1章 PyTorch自动微分核心原理深度解析](#第1章 PyTorch自动微分核心原理深度解析)
-
- [1.1 自动微分（Automatic Differentiation）基础理论](#1.1 自动微分（Automatic Differentiation）基础理论)
-
- [前向模式（Forward Mode）](#前向模式（Forward Mode）)
- [反向模式（Reverse Mode）](#反向模式（Reverse Mode）)
- [1.2 PyTorch Autograd架构设计](#1.2 PyTorch Autograd架构设计)
- [1.3 动态计算图（Dynamic Computation Graph）机制](#1.3 动态计算图（Dynamic Computation Graph）机制)
- [1.4 反向传播引擎实现原理](#1.4 反向传播引擎实现原理)
- [1.5 雅可比向量积（JVP）与向量雅可比积（VJP）](#1.5 雅可比向量积（JVP）与向量雅可比积（VJP）)
- [第2章 Autograd核心组件详解](#第2章 Autograd核心组件详解)
-
- [2.1 Tensor与requires_grad机制](#2.1 Tensor与requires_grad机制)
- [2.2 Function对象与计算历史追踪](#2.2 Function对象与计算历史追踪)
- [2.3 backward()方法深度剖析](#2.3 backward()方法深度剖析)
- [2.4 grad_fn与计算图节点](#2.4 grad_fn与计算图节点)
- [2.5 叶子节点（Leaf Tensor）与梯度累积](#2.5 叶子节点（Leaf Tensor）与梯度累积)
- [第3章 torch.func：可组合函数变换新范式（2024-2025）](#第3章 torch.func：可组合函数变换新范式（2024-2025）)
-
- [3.1 functorch到torch.func的演进](#3.1 functorch到torch.func的演进)
- [3.2 vmap：向量化映射](#3.2 vmap：向量化映射)
- [3.3 grad与grad_and_value](#3.3 grad与grad_and_value)
- [3.4 jvp与vjp：前向与反向模式](#3.4 jvp与vjp：前向与反向模式)
- [3.5 jacrev与jacfwd：雅可比矩阵计算](#3.5 jacrev与jacfwd：雅可比矩阵计算)
- [3.6 hessian：海森矩阵计算](#3.6 hessian：海森矩阵计算)
- [3.7 linearize：线性化变换](#3.7 linearize：线性化变换)
- [第4章 高级特性与实战技巧](#第4章 高级特性与实战技巧)
-
- [4.1 自定义Autograd Function](#4.1 自定义Autograd Function)
- [4.2 梯度检查点（Gradient Checkpointing）](#4.2 梯度检查点（Gradient Checkpointing）)
- [4.3 梯度裁剪与归一化](#4.3 梯度裁剪与归一化)
- [4.4 梯度累积与延迟更新](#4.4 梯度累积与延迟更新)
- [4.5 per-sample gradients实现](#4.5 per-sample gradients实现)
- [4.6 高阶导数计算](#4.6 高阶导数计算)
- 由于篇幅问题，剩下的章节点击这里跳转
-
- [第5章 全链路环境标准化部署](#第5章 全链路环境标准化部署)
-
- [5.1 开发环境配置最佳实践](#5.1 开发环境配置最佳实践)
- [5.2 PyTorch版本兼容性矩阵](#5.2 PyTorch版本兼容性矩阵)
- [5.3 CUDA与cuDNN版本匹配](#5.3 CUDA与cuDNN版本匹配)
- [5.4 分布式训练环境配置](#5.4 分布式训练环境配置)
- [5.5 混合精度训练环境](#5.5 混合精度训练环境)
- [5.6 容器化部署方案](#5.6 容器化部署方案)
- [第6章 实战场景应用案例](#第6章 实战场景应用案例)
-
- [6.1 计算机视觉：CNN梯度优化](#6.1 计算机视觉：CNN梯度优化)
- [6.2 自然语言处理：Transformer梯度流](#6.2 自然语言处理：Transformer梯度流)
- [6.3 图神经网络：消息传递梯度](#6.3 图神经网络：消息传递梯度)
- [6.4 生成对抗网络：双梯度优化](#6.4 生成对抗网络：双梯度优化)
- [6.5 强化学习：策略梯度计算](#6.5 强化学习：策略梯度计算)
- [6.6 物理信息神经网络（PINN）](#6.6 物理信息神经网络（PINN）)
- [第7章 性能优化与调试技巧](#第7章 性能优化与调试技巧)
-
- [7.1 计算图优化策略](#7.1 计算图优化策略)
- [7.2 内存优化技巧](#7.2 内存优化技巧)
- [7.3 梯度流可视化](#7.3 梯度流可视化)
- [7.4 常见梯度问题诊断](#7.4 常见梯度问题诊断)
- [7.5 性能基准测试](#7.5 性能基准测试)
- [第8章 与其他框架对比分析](#第8章 与其他框架对比分析)
-
- [8.1 PyTorch vs TensorFlow自动微分](#8.1 PyTorch vs TensorFlow自动微分)
- [8.2 PyTorch vs JAX自动微分](#8.2 PyTorch vs JAX自动微分)
- [8.3 PyTorch vs MXNet自动微分](#8.3 PyTorch vs MXNet自动微分)
- [8.4 核心差异对比表](#8.4 核心差异对比表)
- [8.5 选型指南与适用场景](#8.5 选型指南与适用场景)
- [第9章 未来发展趋势与演进路线](#第9章 未来发展趋势与演进路线)
-
- [9.1 PyTorch 2.x+新特性展望](#9.1 PyTorch 2.x+新特性展望)
- [9.2 编译优化与torch.compile](#9.2 编译优化与torch.compile)
- [9.3 分布式自动微分](#9.3 分布式自动微分)
- [9.4 量子机器学习集成](#9.4 量子机器学习集成)
- [9.5 自动微分标准化](#9.5 自动微分标准化)
- [第10章 常见问题与解决方案](#第10章 常见问题与解决方案)
- [第11章 总结与最佳实践](#第11章 总结与最佳实践)
- [第12章 附录](#第12章 附录)
📋 目录
摘要
第1章 PyTorch自动微分核心原理深度解析
1.1 自动微分（Automatic Differentiation）基础理论
1.2 PyTorch Autograd架构设计
1.3 动态计算图（Dynamic Computation Graph）机制
1.4 反向传播引擎实现原理
1.5 雅可比向量积（JVP）与向量雅可比积（VJP）
第2章 Autograd核心组件详解
2.1 Tensor与requires_grad机制
2.2 Function对象与计算历史追踪
2.3 backward()方法深度剖析
2.4 grad_fn与计算图节点
2.5 叶子节点（Leaf Tensor）与梯度累积
第3章 torch.func：可组合函数变换新范式（2024-2025）
3.1 functorch到torch.func的演进
3.2 vmap：向量化映射
3.3 grad与grad_and_value
3.4 jvp与vjp：前向与反向模式
3.5 jacrev与jacfwd：雅可比矩阵计算
3.6 hessian：海森矩阵计算
3.7 linearize：线性化变换
第4章 高级特性与实战技巧
4.1 自定义Autograd Function
4.2 梯度检查点（Gradient Checkpointing）
4.3 梯度裁剪与归一化
4.4 梯度累积与延迟更新
4.5 per-sample gradients实现
4.6 高阶导数计算
第5章 全链路环境标准化部署
5.1 开发环境配置最佳实践
5.2 PyTorch版本兼容性矩阵
5.3 CUDA与cuDNN版本匹配
5.4 分布式训练环境配置
5.5 混合精度训练环境
5.6 容器化部署方案
第6章 实战场景应用案例
6.1 计算机视觉：CNN梯度优化
6.2 自然语言处理：Transformer梯度流
6.3 图神经网络：消息传递梯度
6.4 生成对抗网络：双梯度优化
6.5 强化学习：策略梯度计算
6.6 物理信息神经网络（PINN）
第7章 性能优化与调试技巧
7.1 计算图优化策略
7.2 内存优化技巧
7.3 梯度流可视化
7.4 常见梯度问题诊断
7.5 性能基准测试
第8章 与其他框架对比分析
8.1 PyTorch vs TensorFlow自动微分
8.2 PyTorch vs JAX自动微分
8.3 PyTorch vs MXNet自动微分
8.4 核心差异对比表
8.5 选型指南与适用场景
第9章 未来发展趋势与演进路线
9.1 PyTorch 2.x+新特性展望
9.2 编译优化与torch.compile
9.3 分布式自动微分
9.4 量子机器学习集成
9.5 自动微分标准化
第10章 常见问题与解决方案
第11章 总结与最佳实践
附录
附录A：完整代码实现示例
附录B：API参考手册
附录C：性能优化checklist
附录D：术语表与参考文献
📝 摘要
本文介绍了PyTorch开发环境的标准化配置流程和版本兼容性检查。通过自动化脚本实现开发环境的快速搭建，包括系统信息检查、虚拟环境创建、PyTorch安装（根据CUDA版本自动选择）、常用工具包安装等步骤。同时提供了版本兼容性矩阵检查脚本，支持检查PyTorch 2.1.x、2.0.x、1.13.x和1.12.x等版本与Python、CUDA的兼容性关系。该方案可实现PyTorch开发环境的一键标准化部署，确保环境的一致性和可复现性。
本文是一份关于**PyTorch自动微分模块（Autograd）**的全面技术指南，涵盖从底层原理到高级实战的完整知识体系。作为PyTorch的核心引擎，Autograd实现了动态计算图和反向模式自动微分，为深度学习训练提供了强大的基础设施支持。
核心内容：
- ✅ 深度原理剖析：动态计算图构建、反向传播引擎、JVP/VJP机制
- ✅ 最新特性详解：torch.func可组合函数变换（vmap、grad、jvp/vjp等）
- ✅ 高级实战技巧：自定义Function、梯度检查点、per-sample gradients
- ✅ 全链路环境标准化：开发环境、分布式训练、容器化部署
- ✅ 多场景应用案例：CV、NLP、GNN、GAN、RL、PINN等
- ✅ 框架对比选型：PyTorch vs TensorFlow vs JAX深度对比
技术亮点：
- 2024-2025最新特性：全面覆盖torch.func模块的可组合函数变换
- 源码级原理剖析：深入C++后端实现机制
- 工业级最佳实践：环境标准化、性能优化、调试技巧
- 跨框架对比分析：帮助开发者做出正确选型
适用人群：
- 深度学习工程师与研究人员
- PyTorch高级用户
- 自动微分技术爱好者
- 框架选型决策者
第1章 PyTorch自动微分核心原理深度解析
1.1 自动微分（Automatic Differentiation）基础理论
自动微分是深度学习框架的核心技术，它介于符号微分 和数值微分之间：
python
# 三种微分方法对比
import numpy as np
import sympy as sp
import torch
# 1. 符号微分（Symbolic Differentiation）
x = sp.Symbol('x')
f_sym = x**2 + 3*x + 1
df_sym = sp.diff(f_sym, x) # 2*x + 3
print(f"符号微分: {df_sym}")
# 2. 数值微分（Numerical Differentiation）
def f_num(x):
return x**2 + 3*x + 1
def numerical_derivative(x, h=1e-5):
return (f_num(x + h) - f_num(x - h)) / (2 * h)
print(f"数值微分 (x=2): {numerical_derivative(2.0):.6f}")
# 3. 自动微分（Automatic Differentiation）
x = torch.tensor(2.0, requires_grad=True)
f_auto = x**2 + 3*x + 1
f_auto.backward()
print(f"自动微分 (x=2): {x.grad.item():.6f}")自动微分的两种模式：
前向模式（Forward Mode）
python
# 前向模式：从输入到输出，同时计算函数值和导数
# 适合：输入维度 << 输出维度
# 计算复杂度：O(n)，n为输入维度
def forward_mode_ad(x, seed_vector):
"""前向模式自动微分"""
# x: 原始输入
# seed_vector: 种子向量（方向导数方向）
# 前向传播同时计算JVP
y = x**2
dy_dx = 2 * x * seed_vector
return y, dy_dx
x = torch.tensor(3.0)
seed = torch.tensor(1.0)
y, jvp = forward_mode_ad(x, seed)
print(f"前向模式: y={y}, JVP={jvp}")反向模式（Reverse Mode）
python
# 反向模式：先正向计算，再反向传播梯度
# 适合：输入维度 >> 输出维度（深度学习典型场景）
# 计算复杂度：O(m)，m为输出维度
def reverse_mode_ad():
"""反向模式自动微分（PyTorch默认）"""
x = torch.tensor(3.0, requires_grad=True)
y = x**2
# 反向传播
y.backward()
return x.grad
grad = reverse_mode_ad()
print(f"反向模式梯度: {grad}")1.2 PyTorch Autograd架构设计
PyTorch Autograd采用前端-后端分离架构：
┌─────────────────────────────────────────────────┐
│ Python前端层 (torch.autograd) │
├─────────────────────────────────────────────────┤
│ • Tensor.requires_grad │
│ • Tensor.backward() │
│ • torch.autograd.Function │
│ • torch.func.* (vmap, grad, jvp, vjp...) │
└─────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────┐
│ C++后端引擎层 │
├─────────────────────────────────────────────────┤
│ • Engine (反向传播引擎) │
│ • Function (计算图节点) │
│ • Variable (带梯度的张量) │
│ • Node (计算图边) │
└─────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────┐
│ 底层算子库 │
├─────────────────────────────────────────────────┤
│ • ATen (PyTorch张量库) │
│ • THC (CUDA张量库) │
│ • 自定义CUDA扩展 │
└─────────────────────────────────────────────────┘核心数据结构：
cpp
// C++后端核心类（简化版）
class Variable {
Tensor data; // 张量数据
bool requires_grad; // 是否需要梯度
Function* grad_fn; // 梯度函数
std::vector<Variable*> inputs; // 输入变量
};
class Function {
virtual variable_list apply(variable_list&& inputs) = 0; // 前向
virtual variable_list backward(variable_list&& grads) = 0; // 反向
};
class Engine {
void execute(variable_list&& roots, bool keep_graph); // 执行反向传播
};1.3 动态计算图（Dynamic Computation Graph）机制
PyTorch的核心优势在于动态计算图（Define-by-Run）：
python
import torch
from torchviz import make_dot
# 动态计算图示例
def dynamic_computation(x):
"""每次调用都会构建新的计算图"""
h1 = torch.relu(x @ w1 + b1)
# 条件分支 - 动态图的核心优势
if h1.sum() > 0:
h2 = torch.relu(h1 @ w2 + b2)
else:
h2 = torch.sigmoid(h1 @ w3 + b3)
y = h2 @ w4 + b4
return y
# 创建参数
w1 = torch.randn(10, 20, requires_grad=True)
b1 = torch.randn(20, requires_grad=True)
w2 = torch.randn(20, 30, requires_grad=True)
w3 = torch.randn(20, 30, requires_grad=True)
b2 = torch.randn(30, requires_grad=True)
b3 = torch.randn(30, requires_grad=True)
w4 = torch.randn(30, 5, requires_grad=True)
b4 = torch.randn(5, requires_grad=True)
# 第一次调用
x1 = torch.randn(5, 10)
y1 = dynamic_computation(x1)
loss1 = y1.sum()
loss1.backward()
# 第二次调用（可能构建不同的图）
x2 = torch.randn(5, 10)
y2 = dynamic_computation(x2)
loss2 = y2.sum()
loss2.backward()
print("✓ 动态计算图成功构建两次不同的计算路径")与静态图对比：
| 特性 | PyTorch（动态图） | TensorFlow 1.x（静态图） |
|---|---|---|
| 图构建时机 | 运行时 | 定义时 |
| 调试难度 | 容易（Pythonic） | 困难（需Session） |
| 控制流支持 | 原生Python控制流 | 需要特殊API |
| 灵活性 | 高 | 低 |
| 性能优化 | 较难 | 容易 |
1.4 反向传播引擎实现原理
PyTorch反向传播引擎的核心是拓扑排序 和链式法则：
python
# 反向传播引擎简化实现
class BackwardEngine:
def __init__(self):
self.computation_graph = []
def record_operation(self, func, inputs, output):
"""记录前向操作"""
node = {
'func': func,
'inputs': inputs,
'output': output,
'grad_fn': func.backward
}
self.computation_graph.append(node)
def execute_backward(self, output_grad):
"""执行反向传播"""
# 1. 拓扑排序（逆序）
sorted_nodes = reversed(self.computation_graph)
# 2. 初始化梯度字典
gradients = {id(self.computation_graph[-1]['output']): output_grad}
# 3. 按拓扑逆序传播梯度
for node in sorted_nodes:
output_id = id(node['output'])
grad_output = gradients[output_id]
# 4. 调用反向函数
grad_inputs = node['grad_fn'](grad_output, node['inputs'], node['output'])
# 5. 累积梯度
for inp, grad_inp in zip(node['inputs'], grad_inputs):
inp_id = id(inp)
if inp_id in gradients:
gradients[inp_id] = gradients[inp_id] + grad_inp
else:
gradients[inp_id] = grad_inp
return gradients
# PyTorch实际实现（简化）
def actual_pytorch_backward():
"""PyTorch实际的反向传播流程"""
x = torch.tensor(2.0, requires_grad=True)
y = x**2 + 3*x + 1
# 调用.backward()触发反向传播
y.backward()
# 内部流程：
# 1. 从y开始，沿着grad_fn链反向遍历
# 2. 对每个Function节点，调用其backward方法
# 3. 使用链式法则累积梯度到叶子节点
# 4. 清理计算图（除非retain_graph=True）
return x.grad
grad = actual_pytorch_backward()
print(f"PyTorch反向传播结果: {grad}")计算图可视化：
python
# 使用torchviz可视化计算图
import torch
from torchviz import make_dot
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
y = w * x + b
z = y**2
# 生成计算图
graph = make_dot(z, params={'x': x, 'w': w, 'b': b})
graph.render('computation_graph', format='png')
print("✓ 计算图已保存为 computation_graph.png")1.5 雅可比向量积（JVP）与向量雅可比积（VJP）
PyTorch Autograd的核心数学原理是雅可比矩阵的高效计算：
python
import torch
# 雅可比矩阵定义
def jacobian_example():
"""雅可比矩阵示例"""
# 函数: f: R^2 -> R^3
# f(x, y) = [x^2, xy, y^2]
x = torch.tensor([2.0, 3.0], requires_grad=True)
# 计算输出
f1 = x[0]**2 # df1/dx = 2x, df1/dy = 0
f2 = x[0] * x[1] # df2/dx = y, df2/dy = x
f3 = x[1]**2 # df3/dx = 0, df3/dy = 2y
f = torch.stack([f1, f2, f3])
# 计算雅可比矩阵
jacobian = torch.zeros(3, 2)
for i in range(3):
# 对每个输出分量计算梯度
grad = torch.autograd.grad(f[i], x, retain_graph=True)[0]
jacobian[i] = grad
print("雅可比矩阵:")
print(jacobian)
print("\n理论值 (x=2, y=3):")
print("[[4, 0],") # df1/dx=4, df1/dy=0
print(" [3, 2],") # df2/dx=3, df2/dy=2
print(" [0, 6]]") # df3/dx=0, df3/dy=6
return jacobian
jacobian = jacobian_example()
# 向量雅可比积（VJP）- PyTorch默认模式
def vjp_example():
"""向量雅可比积示例"""
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = torch.tensor([1.0, 2.0, 3.0]) # 左乘向量
f1 = x[0]**2
f2 = x[0] * x[1]
f3 = x[1]**2
f = torch.stack([f1, f2, f3])
# VJP: v^T @ J
# 其中v是输出梯度，J是雅可比矩阵
loss = (y * f).sum() # 等价于v^T @ f
loss.backward()
# x.grad = v^T @ J
print(f"VJP结果 (x.grad): {x.grad}")
print(f"理论值: [1*4 + 2*3 + 3*0, 1*0 + 2*2 + 3*6] = [10, 22]")
return x.grad
vjp_result = vjp_example()
# 雅可比向量积（JVP）- torch.func.jvp
def jvp_example():
"""雅可比向量积示例（需要torch.func）"""
try:
from torch.func import jvp
def f(x):
return torch.stack([x[0]**2, x[0]*x[1], x[1]**2])
x = torch.tensor([2.0, 3.0])
v = torch.tensor([1.0, 1.0]) # 右乘向量
# JVP: J @ v
primal, tangent = jvp(f, (x,), (v,))
print(f"JVP结果:")
print(f" 原始输出: {primal}")
print(f" 切向量: {tangent}")
print(f" 理论值: J @ v = [4*1 + 0*1, 3*1 + 2*1, 0*1 + 6*1] = [4, 5, 6]")
return primal, tangent
except ImportError:
print("torch.func.jvp 需要 PyTorch 2.0+")
return None, None
primal, tangent = jvp_example()第2章 Autograd核心组件详解
2.1 Tensor与requires_grad机制
python
import torch
# requires_grad的三种设置方式
# 方式1：创建时设置
x1 = torch.tensor(2.0, requires_grad=True)
# 方式2：使用.requires_grad_()原地修改
x2 = torch.tensor(3.0)
x2.requires_grad_(True)
# 方式3：使用torch.nn.Parameter（推荐用于模型参数）
x3 = torch.nn.Parameter(torch.tensor(4.0))
print(f"x1.requires_grad: {x1.requires_grad}")
print(f"x2.requires_grad: {x2.requires_grad}")
print(f"x3.requires_grad: {x3.requires_grad}")
# requires_grad的传播规则
def requires_grad_propagation():
"""requires_grad传播规则"""
# 规则1：如果所有输入都不需要梯度，输出也不需要
a = torch.tensor(1.0)
b = torch.tensor(2.0)
c = a + b
print(f"规则1 - 无梯度输入: {c.requires_grad}") # False
# 规则2：如果任一输入需要梯度，输出也需要
d = torch.tensor(3.0, requires_grad=True)
e = a + d
print(f"规则2 - 有梯度输入: {e.requires_grad}") # True
# 规则3：使用detach()可以阻断梯度传播
f = d.detach()
g = a + f
print(f"规则3 - detach阻断: {g.requires_grad}") # False
# 规则4：使用torch.no_grad()上下文管理器
with torch.no_grad():
h = a + d
print(f"规则4 - no_grad阻断: {h.requires_grad}") # False
return c, e, g, h
results = requires_grad_propagation()
# 实际应用：冻结部分网络层
def freeze_layers_example():
"""冻结预训练模型的部分层"""
model = torch.nn.Sequential(
torch.nn.Linear(10, 20),
torch.nn.ReLU(),
torch.nn.Linear(20, 30),
torch.nn.ReLU(),
torch.nn.Linear(30, 5)
)
# 冻结前两层
for param in list(model.parameters())[:2]:
param.requires_grad = False
# 只有最后两层会更新
print("可训练参数:")
for name, param in model.named_parameters():
if param.requires_grad:
print(f" {name}: {param.shape}")
return model
frozen_model = freeze_layers_example()2.2 Function对象与计算历史追踪
python
import torch
# 自动追踪的Function对象
def function_tracking_example():
"""Function对象追踪示例"""
x = torch.tensor(2.0, requires_grad=True)
y = x**2 # PowBackward0
z = y + 3 # AddBackward0
w = z * 2 # MulBackward0
print("计算历史追踪:")
print(f"x.grad_fn: {x.grad_fn}") # None (叶子节点)
print(f"y.grad_fn: {y.grad_fn}") # <PowBackward0 object>
print(f"z.grad_fn: {z.grad_fn}") # <AddBackward0 object>
print(f"w.grad_fn: {w.grad_fn}") # <MulBackward0 object>
# 反向传播
w.backward()
print(f"\nx.grad: {x.grad}") # 8.0 = 2 * 2 * 2
# 计算图结构
print("\n计算图结构:")
print("w = 2 * z")
print(" z = y + 3")
print(" y = x^2")
print(" x = 2.0")
return x.grad
grad = function_tracking_example()
# 自定义Function对象
class CustomReLU(torch.autograd.Function):
"""自定义ReLU激活函数"""
@staticmethod
def forward(ctx, input):
"""前向传播"""
# 保存输入用于反向传播
ctx.save_for_backward(input)
# ReLU前向：max(0, x)
output = torch.clamp(input, min=0)
return output
@staticmethod
def backward(ctx, grad_output):
"""反向传播"""
# 获取保存的输入
input, = ctx.saved_tensors
# ReLU反向：grad * (x > 0)
grad_input = grad_output * (input > 0).float()
return grad_input
# 使用自定义Function
def custom_function_usage():
"""自定义Function使用示例"""
x = torch.tensor([-1.0, 0.0, 1.0, 2.0], requires_grad=True)
# 使用自定义ReLU
relu = CustomReLU.apply
y = relu(x)
print(f"输入: {x}")
print(f"输出: {y}")
# 反向传播
loss = y.sum()
loss.backward()
print(f"梯度: {x.grad}")
print(f"理论梯度: [0, 0, 1, 1] (因为负数梯度为0)")
return y, x.grad
output, gradient = custom_function_usage()2.3 backward()方法深度剖析
python
import torch
# backward()的多种用法
def backward_variations():
"""backward()方法的各种用法"""
# 用法1：标量输出（最常见）
print("=== 用法1: 标量输出 ===")
x = torch.tensor(2.0, requires_grad=True)
y = x**2 + 3*x + 1
y.backward() # 等价于 y.backward(torch.tensor(1.0))
print(f"x.grad: {x.grad}") # 7.0
# 用法2：非标量输出 + 自定义梯度
print("\n=== 用法2: 非标量输出 ===")
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x**2 # [1, 4, 9]
# 需要提供gradient参数
gradient = torch.tensor([1.0, 1.0, 1.0])
y.backward(gradient)
print(f"x.grad: {x.grad}") # [2, 4, 6]
# 用法3：retain_graph=True（保留计算图）
print("\n=== 用法3: 保留计算图 ===")
x = torch.tensor(2.0, requires_grad=True)
y = x**2
z = y + 3
z.backward(retain_graph=True)
print(f"第一次反向: x.grad = {x.grad}")
# 可以再次反向传播
z.backward()
print(f"第二次反向: x.grad = {x.grad}") # 累积为16
# 用法4：create_graph=True（创建高阶导数图）
print("\n=== 用法4: 高阶导数 ===")
x = torch.tensor(2.0, requires_grad=True)
y = x**3
# 一阶导数
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
print(f"一阶导数: {dy_dx}") # 12.0
# 二阶导数
d2y_dx2 = torch.autograd.grad(dy_dx, x)[0]
print(f"二阶导数: {d2y_dx2}") # 12.0
return x.grad
result = backward_variations()
# backward()内部机制
def backward_internals():
"""backward()内部机制详解"""
class OperationTracker:
"""操作追踪器"""
def __init__(self):
self.operations = []
def record(self, operation, inputs, output):
self.operations.append({
'operation': operation,
'inputs': inputs,
'output': output
})
# 模拟PyTorch的backward流程
x = torch.tensor(2.0, requires_grad=True)
tracker = OperationTracker()
# 前向传播（记录操作）
y = x**2
tracker.record('pow', [x], y)
z = y + 3
tracker.record('add', [y], z)
w = z * 2
tracker.record('mul', [z], w)
print("前向操作记录:")
for i, op in enumerate(tracker.operations):
print(f" {i+1}. {op['operation']}: {op['inputs']} -> {op['output']}")
# 反向传播
w.backward()
print(f"\n反向传播结果:")
print(f" x.grad = {x.grad}")
print(f" 计算过程: dw/dx = dw/dz * dz/dy * dy/dx")
print(f" = 2 * 1 * 2x = 4x = 8")
return x.grad
grad = backward_internals()2.4 grad_fn与计算图节点
python
import torch
# grad_fn属性详解
def grad_fn_exploration():
"""grad_fn属性探索"""
x = torch.tensor(2.0, requires_grad=True)
y = x**2
z = y + 3
w = z * 2
print("grad_fn属性:")
print(f"x.grad_fn: {x.grad_fn}") # None (叶子节点)
print(f"y.grad_fn: {y.grad_fn}") # <PowBackward0 object>
print(f"z.grad_fn: {z.grad_fn}") # <AddBackward0 object>
print(f"w.grad_fn: {w.grad_fn}") # <MulBackward0 object>
# grad_fn的属性
print("\ngrad_fn详细信息:")
print(f"y.grad_fn.__class__.__name__: {y.grad_fn.__class__.__name__}")
print(f"y.grad_fn.metadata: {y.grad_fn.metadata if hasattr(y.grad_fn, 'metadata') else 'N/A'}")
# 访问grad_fn的输入
print(f"\ny.grad_fn.next_functions:")
for i, (func, _) in enumerate(y.grad_fn.next_functions):
print(f" Input {i}: {func}")
# 构建完整的计算图
def build_computation_graph(tensor, depth=0, visited=None):
"""递归构建计算图"""
if visited is None:
visited = set()
if id(tensor) in visited:
return
visited.add(id(tensor))
indent = " " * depth
print(f"{indent}Tensor: {tensor}")
print(f"{indent} requires_grad: {tensor.requires_grad}")
print(f"{indent} grad_fn: {tensor.grad_fn}")
if tensor.grad_fn:
for next_fn, _ in tensor.grad_fn.next_functions:
if next_fn is not None:
next_tensor = getattr(next_fn, 'variable', None)
if next_tensor is not None:
build_computation_graph(next_tensor, depth + 1, visited)
print("\n完整计算图:")
build_computation_graph(w)
return w
graph_tensor = grad_fn_exploration()
# 自定义grad_fn
class CustomGradFn(torch.autograd.Function):
"""自定义梯度函数"""
@staticmethod
def forward(ctx, x, alpha):
ctx.save_for_backward(x)
ctx.alpha = alpha
return x**2
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
alpha = ctx.alpha
# 自定义梯度：grad * 2 * alpha * x
grad_x = grad_output * 2 * alpha * x
grad_alpha = None # alpha不需要梯度
return grad_x, grad_alpha
@staticmethod
def jvp(ctx, tangent_x, tangent_alpha):
"""前向模式（JVP）"""
x, = ctx.saved_tensors
alpha = ctx.alpha
# JVP: J @ v
# 对于 y = alpha * x^2
# dy/dx = 2*alpha*x, dy/dalpha = x^2
# JVP = [2*alpha*x, x^2] @ [tangent_x, tangent_alpha]
primal_output = ctx.output
tangent_output = 2 * alpha * x * tangent_x + x**2 * tangent_alpha
return primal_output, tangent_output
# 使用自定义grad_fn
def custom_grad_fn_usage():
"""自定义grad_fn使用"""
x = torch.tensor(2.0, requires_grad=True)
alpha = torch.tensor(3.0, requires_grad=True)
custom_op = CustomGradFn.apply
y = custom_op(x, alpha)
print(f"输入: x={x}, alpha={alpha}")
print(f"输出: y={y}")
# 反向传播
y.backward()
print(f"梯度: x.grad={x.grad}, alpha.grad={alpha.grad}")
print(f"理论值: dy/dx=2*alpha*x=12, dy/dalpha=x^2=4")
return y, x.grad, alpha.grad
output, grad_x, grad_alpha = custom_grad_fn_usage()2.5 叶子节点（Leaf Tensor）与梯度累积
python
import torch
# 叶子节点识别
def leaf_tensor_identification():
"""叶子节点识别"""
# 叶子节点：用户创建的、requires_grad=True的张量
x = torch.tensor(2.0, requires_grad=True) # 叶子节点
w = torch.nn.Parameter(torch.tensor(3.0)) # 叶子节点（Parameter）
# 非叶子节点：通过操作生成的张量
y = x**2 # 非叶子节点
z = y + w # 非叶子节点
print("叶子节点识别:")
print(f"x.is_leaf: {x.is_leaf}") # True
print(f"w.is_leaf: {w.is_leaf}") # True
print(f"y.is_leaf: {y.is_leaf}") # False
print(f"z.is_leaf: {z.is_leaf}") # False
# 只有叶子节点会累积梯度
z.backward()
print(f"\n梯度累积:")
print(f"x.grad: {x.grad}") # 4.0
print(f"w.grad: {w.grad}") # 1.0
print(f"y.grad: {y.grad}") # None (非叶子节点默认不保存梯度)
# 保存非叶子节点梯度
y = x**2
y.retain_grad() # 保存非叶子节点梯度
z = y + w
z.backward()
print(f"\n保存非叶子节点梯度:")
print(f"y.grad: {y.grad}") # 1.0
return x.grad, w.grad, y.grad
grads = leaf_tensor_identification()
# 梯度累积机制
def gradient_accumulation():
"""梯度累积机制"""
# 场景：梯度累积用于模拟更大的batch size
model = torch.nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 小batch训练
batch_size = 2
accumulation_steps = 4
for step in range(10):
# 生成小batch数据
x = torch.randn(batch_size, 10)
y_true = torch.randint(0, 5, (batch_size,))
# 前向传播
y_pred = model(x)
loss = torch.nn.functional.cross_entropy(y_pred, y_true)
# 反向传播（梯度累积）
loss.backward()
# 每accumulation_steps步更新一次
if (step + 1) % accumulation_steps == 0:
print(f"Step {step+1}: 更新参数，当前梯度范数: {model.weight.grad.norm():.4f}")
optimizer.step()
optimizer.zero_grad() # 清零梯度
print("\n梯度累积完成")
return model
accumulated_model = gradient_accumulation()
# 梯度清零的重要性
def gradient_zeroing_importance():
"""梯度清零的重要性"""
x = torch.tensor(2.0, requires_grad=True)
optimizer = torch.optim.SGD([x], lr=0.1)
print("不清零梯度的后果:")
for i in range(3):
y = x**2
y.backward()
print(f"Iteration {i+1}:")
print(f" x.grad = {x.grad}")
print(f" x = {x.item()}")
optimizer.step()
# 忘记zero_grad()！
print("\n正确做法（清零梯度）:")
x = torch.tensor(2.0, requires_grad=True)
optimizer = torch.optim.SGD([x], lr=0.1)
for i in range(3):
y = x**2
y.backward()
optimizer.step()
optimizer.zero_grad() # 关键！
print(f"Iteration {i+1}: x = {x.item():.4f}")
return x
result = gradient_zeroing_importance()第3章 torch.func：可组合函数变换新范式（2024-2025）
3.1 functorch到torch.func的演进
python
import torch
# PyTorch 2.0+的torch.func模块
try:
from torch.func import vmap, grad, jvp, vjp, jacrev, jacfwd, hessian, linearize
print("✓ torch.func模块可用")
print("主要功能:")
print(" - vmap: 向量化映射")
print(" - grad: 梯度计算")
print(" - jvp/vjp: 前向/反向模式")
print(" - jacrev/jacfwd: 雅可比矩阵")
print(" - hessian: 海森矩阵")
print(" - linearize: 线性化")
except ImportError:
print("torch.func需要PyTorch 2.0+")
print("安装: pip install torch>=2.0.0")
# functorch (旧版) vs torch.func (新版)
"""
functorch (2021-2022):
- 独立库: pip install functorch
- 导入: from functorch import vmap, grad, ...
torch.func (2023+):
- 内置模块: import torch
- 导入: from torch.func import vmap, grad, ...
- 更好的集成和性能
"""3.2 vmap：向量化映射
python
from torch.func import vmap
import torch
# vmap：将函数映射到批次维度
def vmap_example():
"""vmap示例"""
# 原始函数（处理单个样本）
def single_sample_fn(x):
"""处理单个样本的函数"""
return torch.sin(x) + torch.cos(x)
# 使用vmap向量化
batched_fn = vmap(single_sample_fn)
# 单个样本
x_single = torch.tensor(1.0)
y_single = single_sample_fn(x_single)
print(f"单个样本: x={x_single}, y={y_single}")
# 批次样本（使用vmap）
x_batch = torch.tensor([1.0, 2.0, 3.0])
y_batch = batched_fn(x_batch)
print(f"批次样本: x={x_batch}, y={y_batch}")
# 对比：手动循环
y_manual = torch.stack([single_sample_fn(x) for x in x_batch])
print(f"手动循环结果: {y_manual}")
print(f"vmap与手动循环一致: {torch.allclose(y_batch, y_manual)}")
return y_batch
result = vmap_example()
# vmap在深度学习中的应用
def vmap_in_dl():
"""vmap在深度学习中的应用"""
# 场景1：per-sample gradients（每个样本的梯度）
model = torch.nn.Linear(10, 5)
def loss_fn(params, x, y):
"""计算单个样本的损失"""
y_pred = torch.nn.functional.linear(x, params['weight'], params['bias'])
return torch.nn.functional.cross_entropy(y_pred.unsqueeze(0), y.unsqueeze(0))
# 准备数据
batch_x = torch.randn(32, 10)
batch_y = torch.randint(0, 5, (32,))
# 获取模型参数
params = dict(model.named_parameters())
# 使用vmap计算每个样本的梯度
def compute_per_sample_grad(x, y):
return torch.func.grad(loss_fn)(params, x, y)
# 向量化
per_sample_grads = vmap(compute_per_sample_grad)(batch_x, batch_y)
print("Per-sample gradients:")
print(f" weight梯度形状: {per_sample_grads['weight'].shape}")
print(f" bias梯度形状: {per_sample_grads['bias'].shape}")
# 应用：差分隐私训练
# 对每个样本的梯度进行裁剪和加噪
return per_sample_grads
per_sample_grads = vmap_in_dl()
# vmap多维度映射
def vmap_multi_dimension():
"""vmap多维度映射"""
# 函数：处理二维张量
def matrix_op(a, b):
"""矩阵操作"""
return torch.matmul(a, b)
# 输入：批次的矩阵
batch_a = torch.randn(4, 3, 5) # 4个3x5矩阵
batch_b = torch.randn(4, 5, 2) # 4个5x2矩阵
# 使用vmap在批次维度映射
batched_matmul = vmap(matrix_op)
result = batched_matmul(batch_a, batch_b)
print(f"输入A形状: {batch_a.shape}")
print(f"输入B形状: {batch_b.shape}")
print(f"输出形状: {result.shape}")
print(f"等价于: torch.bmm(batch_a, batch_b)")
# 对比验证
manual_result = torch.bmm(batch_a, batch_b)
print(f"结果一致: {torch.allclose(result, manual_result)}")
return result
result = vmap_multi_dimension()3.3 grad与grad_and_value
python
from torch.func import grad, grad_and_value
import torch
# grad：计算梯度
def grad_example():
"""grad示例"""
# 标量函数
def scalar_fn(x):
"""f(x) = x^2 + 3x + 1"""
return x**2 + 3*x + 1
# 计算梯度
x = torch.tensor(2.0)
gradient_fn = grad(scalar_fn)
grad_x = gradient_fn(x)
print(f"函数: f(x) = x^2 + 3x + 1")
print(f"输入: x = {x}")
print(f"梯度: df/dx = {grad_x}")
print(f"理论值: 2x + 3 = 7")
# 向量函数
def vector_fn(x):
"""f(x) = sum(x^2)"""
return (x**2).sum()
x = torch.tensor([1.0, 2.0, 3.0])
gradient_fn = grad(vector_fn)
grad_x = gradient_fn(x)
print(f"\n函数: f(x) = sum(x^2)")
print(f"输入: x = {x}")
print(f"梯度: df/dx = {grad_x}")
print(f"理论值: [2, 4, 6]")
return grad_x
result = grad_example()
# grad_and_value：同时计算函数值和梯度
def grad_and_value_example():
"""grad_and_value示例"""
def fn(x):
"""f(x) = sin(x) + cos(x)"""
return torch.sin(x) + torch.cos(x)
x = torch.tensor(1.0)
gradient_fn = grad_and_value(fn)
# 同时获取梯度和函数值
grad_x, value = gradient_fn(x)
print(f"函数: f(x) = sin(x) + cos(x)")
print(f"输入: x = {x}")
print(f"函数值: f(x) = {value}")
print(f"梯度: df/dx = {grad_x}")
print(f"理论值: f(1) = sin(1) + cos(1) ≈ 1.3818")
print(f"理论梯度: cos(1) - sin(1) ≈ -0.3012")
# 性能优势：避免重复计算
import time
# 方法1：分别计算
start = time.time()
for _ in range(1000):
val = fn(x)
grad_val = grad(fn)(x)
time_separate = time.time() - start
# 方法2：同时计算
start = time.time()
for _ in range(1000):
grad_val, val = grad_and_value(fn)(x)
time_combined = time.time() - start
print(f"\n性能对比:")
print(f" 分别计算: {time_separate:.4f}s")
print(f" 同时计算: {time_combined:.4f}s")
print(f" 加速比: {time_separate/time_combined:.2f}x")
return grad_x, value
grad_val, fn_val = grad_and_value_example()
# 高阶梯度
def higher_order_gradients():
"""高阶梯度"""
# 二阶导数（海森矩阵的对角线）
def fn(x):
"""f(x) = x^3"""
return x**3
x = torch.tensor(2.0)
# 一阶导数
first_derivative = grad(fn)
df_dx = first_derivative(x)
# 二阶导数
second_derivative = grad(first_derivative)
d2f_dx2 = second_derivative(x)
# 三阶导数
third_derivative = grad(second_derivative)
d3f_dx3 = third_derivative(x)
print(f"函数: f(x) = x^3")
print(f"输入: x = {x}")
print(f"一阶导数: df/dx = {df_dx} (理论: 3x^2 = 12)")
print(f"二阶导数: d²f/dx² = {d2f_dx2} (理论: 6x = 12)")
print(f"三阶导数: d³f/dx³ = {d3f_dx3} (理论: 6)")
# 应用：牛顿法优化
def newton_step(x, lr=0.1):
"""牛顿法更新步骤"""
f = fn(x)
df = first_derivative(x)
d2f = second_derivative(x)
# 牛顿更新: x = x - lr * (df / d2f)
return x - lr * (df / d2f)
# 优化
x = torch.tensor(5.0, requires_grad=True)
for i in range(5):
x = newton_step(x)
print(f"Iteration {i+1}: x = {x.item():.4f}, f(x) = {fn(x).item():.4f}")
return df_dx, d2f_dx2, d3f_dx3
derivatives = higher_order_gradients()3.4 jvp与vjp：前向与反向模式
python
from torch.func import jvp, vjp
import torch
# jvp：雅可比向量积（前向模式）
def jvp_example():
"""jvp示例"""
def fn(x):
"""f(x) = [x^2, sin(x), exp(x)]"""
return torch.stack([x**2, torch.sin(x), torch.exp(x)])
x = torch.tensor(1.0)
v = torch.tensor(1.0) # 切向量
# 计算JVP: J @ v
primal, tangent = jvp(fn, (x,), (v,))
print(f"函数: f(x) = [x^2, sin(x), exp(x)]")
print(f"输入: x = {x}")
print(f"切向量: v = {v}")
print(f"原始输出: {primal}")
print(f"切向量输出: {tangent}")
print(f"理论JVP: J @ v = [2x, cos(x), exp(x)] @ [1] = [2, 0.5403, 2.7183]")
# 多维输入
def multi_input_fn(x, y):
"""f(x, y) = [x^2 + y, x*y]"""
return torch.stack([x**2 + y, x*y])
x = torch.tensor(2.0)
y = torch.tensor(3.0)
v_x = torch.tensor(1.0)
v_y = torch.tensor(0.5)
primal, tangent = jvp(multi_input_fn, (x, y), (v_x, v_y))
print(f"\n函数: f(x, y) = [x^2 + y, x*y]")
print(f"输入: x={x}, y={y}")
print(f"切向量: v_x={v_x}, v_y={v_y}")
print(f"原始输出: {primal}")
print(f"切向量输出: {tangent}")
return primal, tangent
primal, tangent = jvp_example()
# vjp：向量雅可比积（反向模式）
def vjp_example():
"""vjp示例"""
def fn(x):
"""f(x) = [x^2, sin(x), exp(x)]"""
return torch.stack([x**2, torch.sin(x), torch.exp(x)])
x = torch.tensor(1.0)
v = torch.tensor([1.0, 1.0, 1.0]) # 左乘向量
# 计算VJP: v^T @ J
primal, vjp_fn = vjp(fn, x)
# 应用VJP函数
result = vjp_fn(v)
print(f"函数: f(x) = [x^2, sin(x), exp(x)]")
print(f"输入: x = {x}")
print(f"左乘向量: v = {v}")
print(f"原始输出: {primal}")
print(f"VJP结果: {result}")
print(f"理论VJP: v^T @ J = [1, 1, 1] @ [[2x], [cos(x)], [exp(x)]]")
print(f" = 2x + cos(x) + exp(x) = 2 + 0.5403 + 2.7183 = 5.2586")
# 对比：使用标准backward
x = torch.tensor(1.0, requires_grad=True)
y = fn(x)
loss = (y * v).sum() # 等价于v^T @ y
loss.backward()
print(f"\n使用backward验证: x.grad = {x.grad}")
print(f"结果一致: {torch.allclose(result[0], x.grad)}")
return primal, result
primal, vjp_result = vjp_example()
# jvp vs vjp性能对比
def jvp_vs_vjp_performance():
"""jvp vs vjp性能对比"""
# 场景：输入维度高，输出维度低（深度学习典型）
def high_dim_input_fn(x):
"""f: R^1000 -> R^10"""
# 简化：线性变换
w = torch.randn(10, 1000)
return x @ w.T
x = torch.randn(1000)
# JVP：需要运行1000次（输入维度）
import time
start = time.time()
for _ in range(10):
for i in range(1000):
v = torch.zeros(1000)
v[i] = 1.0
_, tangent = jvp(high_dim_input_fn, (x,), (v,))
jvp_time = time.time() - start
# VJP：只需要运行10次（输出维度）
start = time.time()
for _ in range(10):
for i in range(10):
v = torch.zeros(10)
v[i] = 1.0
primal, vjp_fn = vjp(high_dim_input_fn, x)
result = vjp_fn(v)
vjp_time = time.time() - start
print("性能对比 (输入:1000维, 输出:10维):")
print(f" JVP时间: {jvp_time:.4f}s")
print(f" VJP时间: {vjp_time:.4f}s")
print(f" VJP更快: {jvp_time/vjp_time:.2f}x")
print(f"\n结论: 当输入维度 >> 输出维度时，VJP（反向模式）更高效")
return jvp_time, vjp_time
jvp_time, vjp_time = jvp_vs_vjp_performance()3.5 jacrev与jacfwd：雅可比矩阵计算
python
from torch.func import jacrev, jacfwd
import torch
# jacrev：反向模式雅可比矩阵
def jacrev_example():
"""jacrev示例"""
def fn(x):
"""f: R^3 -> R^2"""
# f(x) = [x1^2 + x2, x2*x3 + sin(x1)]
return torch.stack([x[0]**2 + x[1], x[1]*x[2] + torch.sin(x[0])])
x = torch.tensor([1.0, 2.0, 3.0])
# 计算雅可比矩阵
jacobian = jacrev(fn)(x)
print(f"函数: f(x) = [x1^2 + x2, x2*x3 + sin(x1)]")
print(f"输入: x = {x}")
print(f"雅可比矩阵:")
print(jacobian)
print(f"\n理论雅可比:")
print(f" df1/dx1 = 2x1 = 2")
print(f" df1/dx2 = 1")
print(f" df1/dx3 = 0")
print(f" df2/dx1 = cos(x1) = 0.5403")
print(f" df2/dx2 = x3 = 3")
print(f" df2/dx3 = x2 = 2")
print(f"\n矩阵形式:")
print(f" [[2.0000, 1.0000, 0.0000],")
print(f" [0.5403, 3.0000, 2.0000]]")
return jacobian
jacobian = jacrev_example()
# jacfwd：前向模式雅可比矩阵
def jacfwd_example():
"""jacfwd示例"""
def fn(x):
"""f: R^3 -> R^2"""
return torch.stack([x[0]**2 + x[1], x[1]*x[2] + torch.sin(x[0])])
x = torch.tensor([1.0, 2.0, 3.0])
# 计算雅可比矩阵
jacobian = jacfwd(fn)(x)
print(f"使用jacfwd计算的雅可比矩阵:")
print(jacobian)
# 验证与jacrev一致
jacobian_rev = jacrev(fn)(x)
print(f"\njacrev与jacfwd结果一致: {torch.allclose(jacobian, jacobian_rev)}")
return jacobian
jacobian_fwd = jacfwd_example()
# jacrev vs jacfwd性能对比
def jacrev_vs_jacfwd_performance():
"""jacrev vs jacfwd性能对比"""
# 场景1：输入维度 < 输出维度
def small_input_fn(x):
"""f: R^10 -> R^100"""
w = torch.randn(100, 10)
return x @ w.T
x = torch.randn(10)
import time
# jacfwd（前向模式）更优
start = time.time()
for _ in range(10):
J = jacfwd(small_input_fn)(x)
jacfwd_time_small = time.time() - start
# jacrev（反向模式）
start = time.time()
for _ in range(10):
J = jacrev(small_input_fn)(x)
jacrev_time_small = time.time() - start
print("场景1: 输入10维 -> 输出100维")
print(f" jacfwd时间: {jacfwd_time_small:.4f}s")
print(f" jacrev时间: {jacrev_time_small:.4f}s")
print(f" jacfwd更快: {jacrev_time_small/jacfwd_time_small:.2f}x")
# 场景2：输入维度 > 输出维度
def large_input_fn(x):
"""f: R^100 -> R^10"""
w = torch.randn(10, 100)
return x @ w.T
x = torch.randn(100)
start = time.time()
for _ in range(10):
J = jacfwd(large_input_fn)(x)
jacfwd_time_large = time.time() - start
start = time.time()
for _ in range(10):
J = jacrev(large_input_fn)(x)
jacrev_time_large = time.time() - start
print("\n场景2: 输入100维 -> 输出10维")
print(f" jacfwd时间: {jacfwd_time_large:.4f}s")
print(f" jacrev时间: {jacrev_time_large:.4f}s")
print(f" jacrev更快: {jacfwd_time_large/jacrev_time_large:.2f}x")
print("\n结论:")
print(" - 输入维度 < 输出维度: 使用jacfwd（前向模式）")
print(" - 输入维度 > 输出维度: 使用jacrev（反向模式）")
return jacfwd_time_small, jacrev_time_small, jacfwd_time_large, jacrev_time_large
times = jacrev_vs_jacfwd_performance()
# 雅可比矩阵应用：神经网络敏感性分析
def jacobian_in_nn():
"""雅可比矩阵在神经网络中的应用"""
# 简单神经网络
model = torch.nn.Sequential(
torch.nn.Linear(10, 20),
torch.nn.ReLU(),
torch.nn.Linear(20, 5)
)
# 输入样本
x = torch.randn(10)
# 定义预测函数
def predict(params, x):
"""使用给定参数进行预测"""
with torch.no_grad():
# 临时替换模型参数
original_params = dict(model.named_parameters())
for name, param in params.items():
param.data = param.data
output = model(x)
# 恢复原始参数
for name, param in original_params.items():
model.get_parameter(name).data = param.data
return output
# 获取模型参数
params = dict(model.named_parameters())
# 计算输入对输出的雅可比矩阵
# J[i,j] = d(output_i) / d(input_j)
def output_wrt_input(x):
return model(x)
jacobian = jacrev(output_wrt_input)(x)
print(f"输入维度: {x.shape}")
print(f"输出维度: {model(x).shape}")
print(f"雅可比矩阵形状: {jacobian.shape}")
print(f"雅可比矩阵:\n{jacobian}")
# 应用：敏感性分析
# 每个输入特征对每个输出类别的影响
sensitivity = jacobian.abs().mean(dim=0)
print(f"\n输入特征敏感性: {sensitivity}")
return jacobian
jacobian_nn = jacobian_in_nn()3.6 hessian：海森矩阵计算
python
from torch.func import hessian
import torch
# hessian：海森矩阵计算
def hessian_example():
"""hessian示例"""
# 标量函数
def scalar_fn(x):
"""f(x) = x1^2 + 2*x2^2 + 3*x1*x2"""
return x[0]**2 + 2*x[1]**2 + 3*x[0]*x[1]
x = torch.tensor([1.0, 2.0])
# 计算海森矩阵
H = hessian(scalar_fn)(x)
print(f"函数: f(x) = x1^2 + 2*x2^2 + 3*x1*x2")
print(f"输入: x = {x}")
print(f"海森矩阵:")
print(H)
print(f"\n理论海森矩阵:")
print(f" d²f/dx1² = 2")
print(f" d²f/dx1dx2 = 3")
print(f" d²f/dx2dx1 = 3")
print(f" d²f/dx2² = 4")
print(f"\n矩阵形式:")
print(f" [[2, 3],")
print(f" [3, 4]]")
# 特征值分析
eigenvalues = torch.linalg.eigvalsh(H)
print(f"\n特征值: {eigenvalues}")
print(f"正定性: {'是' if (eigenvalues > 0).all() else '否'}")
return H
hessian_matrix = hessian_example()
# 海森矩阵应用：牛顿法优化
def hessian_newton_method():
"""海森矩阵在牛顿法中的应用"""
def loss_fn(x):
"""损失函数: f(x) = (x1 - 3)^2 + 2*(x2 - 2)^2"""
return (x[0] - 3)**2 + 2*(x[1] - 2)**2
# 初始化
x = torch.tensor([0.0, 0.0], requires_grad=True)
print("牛顿法优化:")
print(f"初始点: {x}")
print(f"初始损失: {loss_fn(x):.4f}\n")
for i in range(5):
# 计算梯度和海森矩阵
grad_fn = torch.func.grad(loss_fn)
hess_fn = torch.func.hessian(loss_fn)
g = grad_fn(x)
H = hess_fn(x)
# 牛顿更新: x = x - H^(-1) @ g
delta = torch.linalg.solve(H, g)
x = x - delta
loss = loss_fn(x)
print(f"Iteration {i+1}:")
print(f" x = {x}")
print(f" loss = {loss:.6f}")
print(f" gradient norm = {g.norm():.6f}\n")
if g.norm() < 1e-6:
break
print(f"最优解: x* = {x}")
print(f"最小损失: f(x*) = {loss_fn(x):.6f}")
return x
optimal_x = hessian_newton_method()
# 海森矩阵应用：不确定性量化
def hessian_uncertainty():
"""海森矩阵在不确定性量化中的应用"""
# 假设我们有一个拟合的模型
def model(x, params):
"""线性模型: y = w*x + b"""
w, b = params
return w * x + b
# 生成数据
torch.manual_seed(42)
x_data = torch.linspace(0, 10, 100)
y_true = 2.5 * x_data + 1.5
y_data = y_true + torch.randn(100) * 0.5
# 负对数似然（假设高斯噪声）
def negative_log_likelihood(params):
w, b = params
predictions = w * x_data + b
mse = ((predictions - y_data)**2).mean()
return mse # 简化版
# 找到最优参数
params = torch.tensor([1.0, 1.0], requires_grad=True)
optimizer = torch.optim.LBFGS([params], lr=0.1)
def closure():
optimizer.zero_grad()
loss = negative_log_likelihood(params)
loss.backward()
return loss
for _ in range(10):
optimizer.step(closure)
print(f"最优参数: w={params[0]:.4f}, b={params[1]:.4f}")
# 计算海森矩阵（协方差矩阵的逆）
H = hessian(negative_log_likelihood)(params)
print(f"\n海森矩阵:")
print(H)
# 参数不确定性（协方差矩阵）
covariance = torch.linalg.inv(H)
uncertainties = covariance.diag().sqrt()
print(f"\n参数不确定性:")
print(f" w: {params[0]:.4f} ± {uncertainties[0]:.4f}")
print(f" b: {params[1]:.4f} ± {uncertainties[1]:.4f}")
return params, uncertainties
params, uncertainties = hessian_uncertainty()3.7 linearize：线性化变换
python
from torch.func import linearize
import torch
# linearize：函数线性化
def linearize_example():
"""linearize示例"""
def fn(x):
"""非线性函数: f(x) = sin(x) + x^2"""
return torch.sin(x) + x**2
x0 = torch.tensor(1.0)
# 在x0处线性化
linearized_fn, y0 = linearize(fn, x0)
print(f"函数: f(x) = sin(x) + x^2")
print(f"线性化点: x0 = {x0}")
print(f"函数值: f(x0) = {y0}")
# 线性近似: f(x) ≈ f(x0) + f'(x0) * (x - x0)
def linear_approx(x):
return linearized_fn(x - x0) + y0
# 对比真实值和线性近似
test_points = torch.tensor([0.5, 0.8, 1.0, 1.2, 1.5])
print(f"\n对比真实值和线性近似:")
print(f"{'x':>6} {'f(x)':>10} {'linear(x)':>12} {'error':>10}")
print("-" * 45)
for x in test_points:
true_val = fn(x)
approx_val = linear_approx(x)
error = (true_val - approx_val).abs()
print(f"{x:6.2f} {true_val:10.4f} {approx_val:12.4f} {error:10.4f}")
# 可视化（需要matplotlib）
try:
import matplotlib.pyplot as plt
x_range = torch.linspace(0, 2, 100)
y_true = fn(x_range)
y_linear = linear_approx(x_range)
plt.figure(figsize=(10, 6))
plt.plot(x_range, y_true, label='f(x) = sin(x) + x²', linewidth=2)
plt.plot(x_range, y_linear, '--', label='Linear approximation', linewidth=2)
plt.axvline(x0, color='r', linestyle=':', label=f'Linearization point x0={x0}')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Function Linearization using torch.func.linearize')
plt.legend()
plt.grid(True, alpha=0.3)
plt.savefig('linearization.png', dpi=150, bbox_inches='tight')
print(f"\n✓ 可视化已保存为 linearization.png")
except ImportError:
print("\n提示: 安装matplotlib可查看可视化结果")
return linearized_fn, y0
linearized_fn, y0 = linearize_example()
# linearize在控制系统中的应用
def linearize_control_system():
"""linearize在控制系统中的应用"""
# 非线性动力系统
def dynamics(state, control):
"""
非线性系统动力学
state: [position, velocity]
control: force
"""
pos, vel = state
force = control
# 非线性动力学: 考虑空气阻力
# dv/dt = (force - 0.1*vel^2) / mass
mass = 1.0
drag = 0.1 * vel**2
dpos_dt = vel
dvel_dt = (force - drag) / mass
return torch.stack([dpos_dt, dvel_dt])
# 平衡点
state0 = torch.tensor([0.0, 0.0])
control0 = torch.tensor([0.0])
# 线性化系统
def state_dynamics(state):
return dynamics(state, control0)
A_linearized, _ = linearize(state_dynamics, state0)
def control_dynamics(control):
return dynamics(state0, control)
B_linearized, _ = linearize(control_dynamics, control0)
print("非线性系统线性化:")
print(f"平衡点: state={state0}, control={control0}")
print(f"\n状态矩阵 A (df/dstate):")
print(A_linearized(torch.eye(2)))
print(f"\n控制矩阵 B (df/dcontrol):")
print(B_linearized(torch.eye(1)))
# 线性化系统: dx/dt = A*x + B*u
print(f"\n线性化动力学:")
print(f" dpos/dt = vel")
print(f" dvel/dt = force")
print(f" (在平衡点附近的小扰动)")
return A_linearized, B_linearized
A_lin, B_lin = linearize_control_system()（由于篇幅限制，我将继续撰写后续章节。文章已经涵盖了前3章的详细内容，包括自动微分原理、核心组件和torch.func新特性。接下来的章节将包括高级特性、环境部署、实战案例、性能优化、框架对比等。）
第4章 高级特性与实战技巧
4.1 自定义Autograd Function
python
import torch
# 复杂的自定义Function示例
class ComplexCustomFunction(torch.autograd.Function):
"""复杂的自定义Function，包含多个输入输出"""
@staticmethod
def forward(ctx, x, y, alpha, beta):
"""
前向传播: z = alpha * x^2 + beta * sin(y)
"""
# 保存输入用于反向传播
ctx.save_for_backward(x, y)
ctx.alpha = alpha
ctx.beta = beta
# 计算输出
z1 = alpha * x**2
z2 = beta * torch.sin(y)
z = z1 + z2
# 保存中间结果（可选，用于优化反向传播）
ctx.mark_non_differentiable(alpha, beta)
return z, z1, z2
@staticmethod
def backward(ctx, grad_z, grad_z1, grad_z2):
"""
反向传播
"""
# 获取保存的变量
x, y = ctx.saved_tensors
alpha = ctx.alpha
beta = ctx.beta
# 计算梯度
# dz/dx = 2 * alpha * x
# dz/dy = beta * cos(y)
grad_x = grad_z * 2 * alpha * x
grad_y = grad_z * beta * torch.cos(y)
# alpha和beta不需要梯度（已标记为non-differentiable）
grad_alpha = None
grad_beta = None
return grad_x, grad_y, grad_alpha, grad_beta
@staticmethod
def jvp(ctx, tangent_x, tangent_y, tangent_alpha, tangent_beta):
"""
前向模式（JVP）
"""
x, y = ctx.saved_tensors
alpha = ctx.alpha
beta = ctx.beta
# 原始输出
primal_z, primal_z1, primal_z2 = ctx.output
# JVP计算
# dz = 2*alpha*x*dx + x^2*dalpha + beta*cos(y)*dy + sin(y)*dbeta
tangent_z = (2 * alpha * x * tangent_x +
x**2 * tangent_alpha +
beta * torch.cos(y) * tangent_y +
torch.sin(y) * tangent_beta)
tangent_z1 = 2 * alpha * x * tangent_x + x**2 * tangent_alpha
tangent_z2 = beta * torch.cos(y) * tangent_y + torch.sin(y) * tangent_beta
return (primal_z, primal_z1, primal_z2), (tangent_z, tangent_z1, tangent_z2)
# 使用复杂自定义Function
def complex_custom_function_usage():
"""复杂自定义Function使用示例"""
x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(1.0, requires_grad=True)
alpha = torch.tensor(3.0)
beta = torch.tensor(2.0)
custom_op = ComplexCustomFunction.apply
z, z1, z2 = custom_op(x, y, alpha, beta)
print(f"输入: x={x}, y={y}, alpha={alpha}, beta={beta}")
print(f"输出: z={z}, z1={z1}, z2={z2}")
# 反向传播
loss = z.sum()
loss.backward()
print(f"\n梯度:")
print(f" x.grad = {x.grad}")
print(f" y.grad = {y.grad}")
print(f" alpha.grad = {alpha.grad}") # None
print(f" beta.grad = {beta.grad}") # None
# 理论验证
print(f"\n理论梯度:")
print(f" dz/dx = 2*alpha*x = {2*3*2}")
print(f" dz/dy = beta*cos(y) = {2*torch.cos(torch.tensor(1.0))}")
return z, x.grad, y.grad
output, grad_x, grad_y = complex_custom_function_usage()
# 自定义Function在CUDA上的实现
class CUDACustomFunction(torch.autograd.Function):
"""支持CUDA的自定义Function"""
@staticmethod
def forward(ctx, x):
# 保存输入
ctx.save_for_backward(x)
# 根据设备选择实现
if x.is_cuda:
# CUDA实现（简化）
return x**2
else:
# CPU实现
return x**2
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
if x.is_cuda:
# CUDA反向
return 2 * x * grad_output
else:
# CPU反向
return 2 * x * grad_output
# CUDA自定义Function使用
def cuda_custom_function_usage():
"""CUDA自定义Function使用"""
# CPU版本
x_cpu = torch.tensor(2.0, requires_grad=True)
y_cpu = CUDACustomFunction.apply(x_cpu)
y_cpu.backward()
print(f"CPU版本: x.grad = {x_cpu.grad}")
# CUDA版本（如果有GPU）
if torch.cuda.is_available():
x_cuda = torch.tensor(2.0, requires_grad=True, device='cuda')
y_cuda = CUDACustomFunction.apply(x_cuda)
y_cuda.backward()
print(f"CUDA版本: x.grad = {x_cuda.grad}")
print(f"结果一致: {torch.allclose(x_cpu.grad, x_cuda.grad.cpu())}")
return x_cpu.grad
grad = cuda_custom_function_usage()方案二
python
import torch
from torch.autograd import Function
# 基础自定义Autograd Function
class CustomReLU(Function):
"""自定义ReLU函数及其梯度"""
@staticmethod
def forward(ctx, input):
"""
前向传播
Args:
ctx: 上下文对象，用于保存信息供反向传播使用
input: 输入张量
Returns:
输出张量
"""
# 保存输入张量供反向传播使用
ctx.save_for_backward(input)
# ReLU前向计算: max(0, x)
output = torch.clamp(input, min=0)
return output
@staticmethod
def backward(ctx, grad_output):
"""
反向传播
Args:
ctx: 上下文对象
grad_output: 输出梯度
Returns:
输入梯度
"""
# 获取保存的输入
input, = ctx.saved_tensors
# ReLU梯度: x > 0 时为1，否则为0
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input
# 使用自定义函数
def basic_custom_autograd():
"""基础自定义Autograd示例"""
# 创建输入（需要梯度）
x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0], requires_grad=True)
# 使用自定义ReLU
output = CustomReLU.apply(x)
print("自定义ReLU示例:")
print(f"输入: {x}")
print(f"输出: {output}")
# 计算梯度
loss = output.sum()
loss.backward()
print(f"梯度: {x.grad}")
print(f"\n验证:")
print(f" 输入 < 0 的梯度应为 0: {x.grad[x < 0]}")
print(f" 输入 >= 0 的梯度应为 1: {x.grad[x >= 0]}")
return output, x.grad
output, grad = basic_custom_autograd()
# 复杂自定义函数：带参数的自定义操作
class CustomDropout(Function):
"""自定义Dropout函数"""
@staticmethod
def forward(ctx, input, p=0.5, training=True):
"""
Dropout前向传播
Args:
ctx: 上下文
input: 输入张量
p: dropout概率
training: 是否训练模式
Returns:
dropout后的输出
"""
ctx.p = p
ctx.training = training
if not training:
return input
# 生成mask
mask = torch.bernoulli(torch.ones_like(input) * (1 - p))
# 保存mask供反向传播
ctx.save_for_backward(mask)
# 应用dropout并缩放
output = input * mask / (1 - p)
return output
@staticmethod
def backward(ctx, grad_output):
"""Dropout反向传播"""
if not ctx.training:
return grad_output, None, None
mask, = ctx.saved_tensors
# 只有被保留的神经元才有梯度
grad_input = grad_output * mask / (1 - ctx.p)
return grad_input, None, None
# 使用自定义Dropout
def custom_dropout_example():
"""自定义Dropout示例"""
x = torch.randn(3, 4, requires_grad=True)
print("自定义Dropout示例:")
print(f"输入:\n{x}\n")
# 训练模式
output_train = CustomDropout.apply(x, p=0.5, training=True)
print(f"训练模式输出:\n{output_train}\n")
loss = output_train.sum()
loss.backward()
print(f"训练模式梯度:\n{x.grad}\n")
# 评估模式
x.grad.zero_()
output_eval = CustomDropout.apply(x, p=0.5, training=False)
print(f"评估模式输出:\n{output_eval}")
print(f"（评估模式下输出等于输入）")
return output_train, output_eval
train_out, eval_out = custom_dropout_example()
# 高级示例：自定义数值稳定操作
class StableSoftmax(Function):
"""数值稳定的Softmax实现"""
@staticmethod
def forward(ctx, input, dim=-1):
"""
Softmax前向传播
Args:
ctx: 上下文
input: 输入张量
dim: 计算维度
Returns:
softmax输出
"""
ctx.dim = dim
# 数值稳定技巧：减去最大值
max_val = input.max(dim=dim, keepdim=True)[0]
input_stable = input - max_val
# 计算exp
exp_input = torch.exp(input_stable)
# 计算sum
sum_exp = exp_input.sum(dim=dim, keepdim=True)
# 计算softmax
output = exp_input / sum_exp
# 保存信息供反向传播
ctx.save_for_backward(output)
return output
@staticmethod
def backward(ctx, grad_output):
"""Softmax反向传播"""
output, = ctx.saved_tensors
dim = ctx.dim
# Softmax梯度公式: grad * (output - output^2)
# 更高效的形式: grad_output - output * sum(grad_output * output)
grad_input = grad_output - output * torch.sum(
grad_output * output,
dim=ctx.dim,
keepdim=True
)
return grad_input, None
# 使用自定义Softmax
def stable_softmax_example():
"""数值稳定Softmax示例"""
# 测试数值稳定性
x = torch.tensor([[1000.0, 1000.0, 1000.0]], requires_grad=True)
print("数值稳定Softmax示例:")
print(f"输入: {x}")
# 使用自定义Softmax
output = StableSoftmax.apply(x, dim=-1)
print(f"输出: {output}")
# 验证梯度
loss = output.sum()
loss.backward()
print(f"梯度: {x.grad}")
# 对比PyTorch内置Softmax
x2 = torch.tensor([[1000.0, 1000.0, 1000.0]], requires_grad=True)
output_builtin = torch.softmax(x2, dim=-1)
print(f"\nPyTorch内置Softmax输出: {output_builtin}")
# 验证两者结果一致
print(f"\n结果一致: {torch.allclose(output, output_builtin)}")
return output
output = stable_softmax_example()
# 自定义函数的高级特性：多次保存和恢复
class MultiSaveFunction(Function):
"""演示多次保存和恢复"""
@staticmethod
def forward(ctx, x, y):
# 保存多个张量
ctx.save_for_backward(x, y)
# 也可以保存非张量对象
ctx.constant = 2.5
return x * y + ctx.constant
@staticmethod
def backward(ctx, grad_output):
# 恢复保存的张量
x, y = ctx.saved_tensors
constant = ctx.constant
# 计算梯度
grad_x = grad_output * y
grad_y = grad_output * x
# 注意：常量没有梯度
return grad_x, grad_y
# 使用示例
def multi_save_example():
"""多次保存示例"""
x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)
output = MultiSaveFunction.apply(x, y)
print(f"MultiSave示例:")
print(f" x={x.item()}, y={y.item()}")
print(f" 输出: {output.item()}")
print(f" 预期: 2*3 + 2.5 = 8.5")
output.backward()
print(f" x梯度: {x.grad.item()} (预期: 3)")
print(f" y梯度: {y.grad.item()} (预期: 2)")
return output
output = multi_save_example()
# 自定义函数在实际模型中的应用
class CustomLayer(torch.nn.Module):
"""使用自定义Autograd的自定义层"""
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.randn(10, 10))
self.bias = torch.nn.Parameter(torch.randn(10))
def forward(self, x):
# 使用自定义函数进行线性变换
return CustomLinear.apply(x, self.weight, self.bias)
class CustomLinear(Function):
"""自定义线性变换"""
@staticmethod
def forward(ctx, input, weight, bias):
ctx.save_for_backward(input, weight, bias)
output = input.matmul(weight.t()) + bias
return output
@staticmethod
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_tensors
grad_input = grad_output.matmul(weight)
grad_weight = grad_output.t().matmul(input)
grad_bias = grad_output.sum(0)
return grad_input, grad_weight, grad_bias
# 使用自定义层
def custom_layer_example():
"""自定义层示例"""
layer = CustomLayer()
x = torch.randn(5, 10, requires_grad=True)
output = layer(x)
loss = output.sum()
loss.backward()
print("自定义层示例:")
print(f"输入形状: {x.shape}")
print(f"输出形状: {output.shape}")
print(f"权重梯度形状: {layer.weight.grad.shape}")
print(f"偏置梯度形状: {layer.bias.grad.shape}")
return layer
custom_layer = custom_layer_example()
# 性能优化：使用inplace操作
class InplaceReLU(Function):
"""使用inplace操作的ReLU"""
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
# inplace操作
output = input.clamp_(min=0)
return output
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input
# 对比性能
def inplace_performance_comparison():
"""inplace操作性能对比"""
import time
x = torch.randn(1000, 1000, requires_grad=True)
# 非inplace版本
class RegularReLU(Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return torch.clamp(input, min=0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input
# 性能测试
num_iterations = 1000
# Regular ReLU
start = time.time()
for _ in range(num_iterations):
x.grad = None
output = RegularReLU.apply(x.clone())
output.sum().backward()
regular_time = time.time() - start
# Inplace ReLU
start = time.time()
for _ in range(num_iterations):
x.grad = None
output = InplaceReLU.apply(x.clone())
output.sum().backward()
inplace_time = time.time() - start
print("Inplace操作性能对比:")
print(f"Regular ReLU时间: {regular_time:.4f}s")
print(f"Inplace ReLU时间: {inplace_time:.4f}s")
print(f"加速比: {regular_time/inplace_time:.2f}x")
return regular_time, inplace_time
regular_time, inplace_time = inplace_performance_comparison()4.2 梯度检查点（Gradient Checkpointing）
梯度检查点是一种内存优化技术，通过在前向传播时丢弃中间激活值，在反向传播时重新计算来节省显存。
python
import torch
import torch.utils.checkpoint as checkpoint
# 基础梯度检查点示例
def basic_gradient_checkpointing():
"""基础梯度检查点示例"""
# 定义一个深层网络
class DeepNetwork(torch.nn.Module):
def __init__(self, num_layers=10):
super().__init__()
self.layers = torch.nn.ModuleList([
torch.nn.Linear(1024, 1024) for _ in range(num_layers)
])
def forward(self, x):
for layer in self.layers:
x = torch.nn.functional.relu(layer(x))
return x
# 不使用梯度检查点
def forward_without_checkpoint(model, x):
return model(x)
# 使用梯度检查点
def forward_with_checkpoint(model, x):
def custom_forward(*inputs):
x = inputs[0]
for layer in model.layers:
x = torch.nn.functional.relu(layer(x))
return x
return checkpoint.checkpoint(custom_forward, x)
# 对比显存使用
model = DeepNetwork(num_layers=20)
x = torch.randn(32, 1024, requires_grad=True)
print("梯度检查点对比:")
print(f"输入形状: {x.shape}")
print(f"模型层数: 20")
# 不使用检查点
torch.cuda.reset_peak_memory_stats()
output1 = forward_without_checkpoint(model, x)
output1.sum().backward()
mem_without = torch.cuda.max_memory_allocated() / 1024**2
print(f"不使用检查点 - 显存: {mem_without:.2f} MB")
# 使用检查点
model.zero_grad()
torch.cuda.reset_peak_memory_stats()
output2 = forward_with_checkpoint(model, x)
output2.sum().backward()
mem_with = torch.cuda.max_memory_allocated() / 1024**2
print(f"使用检查点 - 显存: {mem_with:.2f} MB")
print(f"显存节省: {(mem_without - mem_with) / mem_without * 100:.1f}%")
return output1, output2
output1, output2 = basic_gradient_checkpointing()
# 序列模型的梯度检查点
def sequence_model_checkpointing():
"""序列模型的梯度检查点"""
class TransformerBlock(torch.nn.Module):
def __init__(self, dim=512, num_heads=8):
super().__init__()
self.attn = torch.nn.MultiheadAttention(dim, num_heads, batch_first=True)
self.ffn = torch.nn.Sequential(
torch.nn.Linear(dim, dim * 4),
torch.nn.GELU(),
torch.nn.Linear(dim * 4, dim)
)
self.norm1 = torch.nn.LayerNorm(dim)
self.norm2 = torch.nn.LayerNorm(dim)
def forward(self, x):
# 自注意力
attn_out, _ = self.attn(x, x, x)
x = x + attn_out
x = self.norm1(x)
# 前馈网络
ffn_out = self.ffn(x)
x = x + ffn_out
x = self.norm2(x)
return x
class CheckpointedTransformer(torch.nn.Module):
def __init__(self, num_layers=12, dim=512):
super().__init__()
self.layers = torch.nn.ModuleList([
TransformerBlock(dim=dim) for _ in range(num_layers)
])
def forward(self, x):
# 对每个Transformer层使用检查点
for layer in self.layers:
x = checkpoint.checkpoint(layer, x)
return x
# 使用示例
model = CheckpointedTransformer(num_layers=24, dim=1024)
x = torch.randn(8, 128, 1024, requires_grad=True) # batch_size=8, seq_len=128, dim=1024
print("Transformer梯度检查点:")
print(f"输入形状: {x.shape}")
print(f"层数: 24")
output = model(x)
loss = output.sum()
loss.backward()
print(f"输出形状: {output.shape}")
print(f"✓ 梯度检查点成功应用")
return model, output
model, output = sequence_model_checkpointing()
# 高级梯度检查点：分段检查点
def segmented_checkpointing():
"""分段梯度检查点 - 平衡显存和计算"""
class SegmentedModel(torch.nn.Module):
def __init__(self, num_layers=24):
super().__init__()
self.layers = torch.nn.ModuleList([
torch.nn.Linear(1024, 1024) for _ in range(num_layers)
])
def forward(self, x):
# 将24层分为3个段，每段8层
segment_size = 8
for i in range(0, len(self.layers), segment_size):
segment_layers = self.layers[i:i+segment_size]
def create_segment_forward(layers):
def segment_forward(x):
for layer in layers:
x = torch.nn.functional.relu(layer(x))
return x
return segment_forward
x = checkpoint.checkpoint(
create_segment_forward(segment_layers),
x,
use_reentrant=False # PyTorch 2.0+推荐
)
return x
model = SegmentedModel(num_layers=24)
x = torch.randn(16, 1024, requires_grad=True)
print("分段梯度检查点:")
print(f"总层数: 24")
print(f"分段大小: 8层/段")
print(f"段数: 3")
output = model(x)
output.sum().backward()
print(f"✓ 分段检查点成功")
return model, output
model, output = segmented_checkpointing()
# 梯度检查点与混合精度结合
def checkpoint_with_amp():
"""梯度检查点与混合精度训练结合"""
from torch.cuda.amp import autocast, GradScaler
class AMPModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.layers = torch.nn.Sequential(
torch.nn.Linear(1024, 2048),
torch.nn.ReLU(),
torch.nn.Linear(2048, 2048),
torch.nn.ReLU(),
torch.nn.Linear(2048, 1024)
)
def forward(self, x):
return self.layers(x)
model = AMPModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
# 使用梯度检查点的训练步骤
def train_step(x, y):
optimizer.zero_grad()
with autocast():
def forward_function(x):
return model(x)
output = checkpoint.checkpoint(forward_function, x)
loss = torch.nn.functional.mse_loss(output, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
return loss
# 训练循环
for step in range(10):
x = torch.randn(64, 1024).cuda()
y = torch.randn(64, 1024).cuda()
loss = train_step(x, y)
if step % 3 == 0:
print(f"Step {step}: Loss = {loss.item():.4f}")
print("✓ AMP + Checkpoint训练完成")
return model
trained_model = checkpoint_with_amp()4.3 梯度裁剪与归一化
梯度裁剪是防止梯度爆炸的重要技术，特别是在训练RNN和Transformer时。
python
import torch
import torch.nn as nn
# 梯度裁剪方法对比
def gradient_clipping_methods():
"""梯度裁剪方法对比"""
model = nn.Sequential(
nn.Linear(100, 200),
nn.ReLU(),
nn.Linear(200, 100)
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# 生成数据
x = torch.randn(32, 100)
y = torch.randn(32, 100)
# 方法1: 梯度范数裁剪（最常用）
def clip_grad_norm_example():
optimizer.zero_grad()
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss.backward()
# 裁剪梯度范数
max_norm = 1.0
total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm)
print(f"方法1 - 梯度范数裁剪:")
print(f" 裁剪前梯度范数: {total_norm:.4f}")
print(f" 裁剪阈值: {max_norm}")
optimizer.step()
# 方法2: 梯度值裁剪
def clip_grad_value_example():
optimizer.zero_grad()
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss.backward()
# 裁剪梯度值
clip_value = 0.5
nn.utils.clip_grad_value_(model.parameters(), clip_value)
print(f"\n方法2 - 梯度值裁剪:")
print(f" 裁剪阈值: ±{clip_value}")
optimizer.step()
# 方法3: 自定义裁剪策略
def custom_clipping():
optimizer.zero_grad()
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss.backward()
# 按层裁剪
for name, param in model.named_parameters():
if param.grad is not None:
# 不同层使用不同的裁剪阈值
if '0' in name: # 第一层
max_norm = 0.5
else: # 其他层
max_norm = 1.0
param_norm = param.grad.data.norm(2)
clip_coef = max_norm / (param_norm + 1e-6)
if clip_coef < 1:
param.grad.data.mul_(clip_coef)
print(f"\n方法3 - 自定义按层裁剪:")
print(f" 第一层阈值: 0.5")
print(f" 其他层阈值: 1.0")
optimizer.step()
clip_grad_norm_example()
clip_grad_value_example()
custom_clipping()
return model
model = gradient_clipping_methods()
# 梯度归一化与梯度累积结合
def gradient_normalization_with_accumulation():
"""梯度归一化与梯度累积结合"""
class LargeModel(nn.Module):
def __init__(self):
super().__init__()
self.encoder = nn.Linear(512, 1024)
self.decoder = nn.Linear(1024, 512)
def forward(self, x):
x = self.encoder(x)
x = nn.functional.relu(x)
x = self.decoder(x)
return x
model = LargeModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
batch_size = 8
accumulation_steps = 4
max_grad_norm = 1.0
print("梯度累积 + 梯度裁剪:")
print(f"实际batch_size: {batch_size}")
print(f"累积步数: {accumulation_steps}")
print(f"有效batch_size: {batch_size * accumulation_steps}")
print(f"梯度裁剪阈值: {max_grad_norm}\n")
for step in range(10):
# 模拟小batch训练
x = torch.randn(batch_size, 512)
y = torch.randn(batch_size, 512)
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss = loss / accumulation_steps # 归一化损失
loss.backward()
# 每accumulation_steps步更新一次
if (step + 1) % accumulation_steps == 0:
# 裁剪累积的梯度
total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
print(f"Step {step+1}:")
print(f" Loss: {loss.item() * accumulation_steps:.4f}")
print(f" Gradient norm: {total_norm:.4f}")
optimizer.step()
optimizer.zero_grad()
print("\n✓ 梯度累积与裁剪完成")
return model
model = gradient_normalization_with_accumulation()
# 动态梯度裁剪
def dynamic_gradient_clipping():
"""动态梯度裁剪 - 根据训练阶段调整"""
class DynamicClipper:
def __init__(self, initial_max_norm=5.0, min_max_norm=0.5, decay_rate=0.99):
self.max_norm = initial_max_norm
self.min_max_norm = min_max_norm
self.decay_rate = decay_rate
self.step = 0
def clip(self, parameters):
"""执行裁剪并更新阈值"""
total_norm = nn.utils.clip_grad_norm_(parameters, self.max_norm)
# 动态调整裁剪阈值
self.step += 1
self.max_norm = max(
self.min_max_norm,
self.max_norm * (self.decay_rate ** self.step)
)
return total_norm, self.max_norm
model = nn.Sequential(
nn.Linear(100, 200),
nn.ReLU(),
nn.Linear(200, 100)
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
clipper = DynamicClipper(initial_max_norm=5.0, min_max_norm=0.5, decay_rate=0.995)
print("动态梯度裁剪:")
print(f"初始阈值: 5.0")
print(f"最小阈值: 0.5")
print(f"衰减率: 0.995\n")
for step in range(20):
x = torch.randn(32, 100)
y = torch.randn(32, 100)
optimizer.zero_grad()
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss.backward()
grad_norm, current_threshold = clipper.clip(model.parameters())
optimizer.step()
if step % 5 == 0:
print(f"Step {step}: Loss={loss.item():.4f}, "
f"GradNorm={grad_norm:.4f}, "
f"Threshold={current_threshold:.4f}")
print("\n✓ 动态梯度裁剪完成")
return model, clipper
model, clipper = dynamic_gradient_clipping()4.4 梯度累积与延迟更新
梯度累积允许使用更小的batch size模拟更大的batch size，特别适用于显存受限的场景。
python
import torch
import torch.nn as nn
# 基础梯度累积
def basic_gradient_accumulation():
"""基础梯度累积"""
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(100, 200)
self.fc2 = nn.Linear(200, 100)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# 配置
actual_batch_size = 8
accumulation_steps = 4
effective_batch_size = actual_batch_size * accumulation_steps
print("梯度累积示例:")
print(f"实际batch_size: {actual_batch_size}")
print(f"累积步数: {accumulation_steps}")
print(f"有效batch_size: {effective_batch_size}\n")
# 训练循环
for epoch in range(3):
epoch_loss = 0
num_batches = 0
for step in range(20): # 模拟20个batch
# 生成数据
x = torch.randn(actual_batch_size, 100)
y = torch.randn(actual_batch_size, 100)
# 前向传播
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss = loss / accumulation_steps # 归一化
# 反向传播
loss.backward()
epoch_loss += loss.item()
num_batches += 1
# 每accumulation_steps步更新一次
if (step + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
avg_loss = epoch_loss / num_batches
print(f"Epoch {epoch+1}: Average Loss = {avg_loss:.6f}")
print("\n✓ 梯度累积训练完成")
return model
model = basic_gradient_accumulation()
# 梯度累积与学习率调整
def gradient_accumulation_with_lr_schedule():
"""梯度累积与学习率调整"""
class LRAdjustedOptimizer:
def __init__(self, optimizer, accumulation_steps):
self.optimizer = optimizer
self.accumulation_steps = accumulation_steps
self.step_count = 0
def step(self):
self.step_count += 1
# 只在累积完成时更新
if self.step_count % self.accumulation_steps == 0:
self.optimizer.step()
self.optimizer.zero_grad()
def zero_grad(self):
if self.step_count % self.accumulation_steps == 0:
self.optimizer.zero_grad()
model = nn.Sequential(
nn.Linear(100, 200),
nn.ReLU(),
nn.Linear(200, 100)
)
# 基础学习率需要根据有效batch size调整
base_lr = 1e-3
accumulation_steps = 4
# 线性缩放规则：lr ∝ batch_size
adjusted_lr = base_lr * accumulation_steps
optimizer = torch.optim.Adam(model.parameters(), lr=adjusted_lr)
wrapped_optimizer = LRAdjustedOptimizer(optimizer, accumulation_steps)
print("梯度累积与学习率调整:")
print(f"基础学习率: {base_lr}")
print(f"累积步数: {accumulation_steps}")
print(f"调整后学习率: {adjusted_lr}")
print(f"缩放因子: {accumulation_steps}x\n")
for step in range(20):
x = torch.randn(8, 100)
y = torch.randn(8, 100)
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss.backward()
wrapped_optimizer.step()
if step % 5 == 0:
print(f"Step {step}: Loss = {loss.item():.6f}")
print("\n✓ 学习率调整的梯度累积完成")
return model, wrapped_optimizer
model, optimizer = gradient_accumulation_with_lr_schedule()
# 梯度累积与混合精度训练
def gradient_accumulation_with_amp():
"""梯度累积与混合精度训练"""
from torch.cuda.amp import autocast, GradScaler
if not torch.cuda.is_available():
print("CUDA不可用，跳过AMP示例")
return None
class AMPModel(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(512, 1024),
nn.ReLU(),
nn.Linear(1024, 1024),
nn.ReLU(),
nn.Linear(1024, 512)
)
def forward(self, x):
return self.layers(x)
model = AMPModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
accumulation_steps = 4
print("梯度累积 + AMP:")
print(f"累积步数: {accumulation_steps}\n")
for step in range(20):
x = torch.randn(16, 512).cuda()
y = torch.randn(16, 512).cuda()
# 混合精度前向
with autocast():
output = model(x)
loss = nn.functional.mse_loss(output, y)
loss = loss / accumulation_steps
# 缩放损失并反向传播
scaler.scale(loss).backward()
# 每accumulation_steps步更新
if (step + 1) % accumulation_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
if step % 5 == 0:
print(f"Step {step}: Loss = {loss.item() * accumulation_steps:.6f}")
print("\n✓ AMP + 梯度累积完成")
return model
model = gradient_accumulation_with_amp()4.5 per-sample gradients实现
per-sample gradients是差分隐私训练和某些高级优化算法的关键技术。
python
import torch
import torch.nn as nn
from torch.func import vmap, grad
# per-sample gradients基础实现
def per_sample_gradients_basic():
"""per-sample gradients基础实现"""
# 方法1: 使用循环（慢但直观）
def per_sample_grads_loop(model, x_batch, y_batch):
"""使用循环计算每个样本的梯度"""
per_sample_grads = []
for i in range(len(x_batch)):
x = x_batch[i:i+1]
y = y_batch[i:i+1]
# 前向传播
output = model(x)
loss = nn.functional.cross_entropy(output, y)
# 计算梯度
model.zero_grad()
loss.backward()
# 收集梯度
grads = {}
for name, param in model.named_parameters():
if param.grad is not None:
grads[name] = param.grad.clone()
per_sample_grads.append(grads)
return per_sample_grads
# 方法2: 使用torch.func.vmap（推荐，高效）
def per_sample_grads_vmap(model, x_batch, y_batch):
"""使用vmap计算每个样本的梯度"""
def compute_loss(params, x, y):
"""计算单个样本的损失"""
output = torch.functional.functional_call(model, params, x)
return nn.functional.cross_entropy(output, y)
# 获取模型参数
params = dict(model.named_parameters())
# 使用vmap向量化梯度计算
def compute_grad(x, y):
return grad(compute_loss)(params, x.unsqueeze(0), y.unsqueeze(0))
per_sample_grads = vmap(compute_grad)(x_batch, y_batch)
return per_sample_grads
# 创建模型和数据
model = nn.Linear(10, 5)
x_batch = torch.randn(32, 10)
y_batch = torch.randint(0, 5, (32,))
print("Per-sample gradients对比:")
print(f"Batch size: {len(x_batch)}")
print(f"输入形状: {x_batch.shape}")
print(f"输出类别数: 5\n")
# 使用vmap方法（高效）
per_sample_grads = per_sample_grads_vmap(model, x_batch, y_batch)
print("使用torch.func.vmap:")
for name, grads in per_sample_grads.items():
print(f" {name}: {grads.shape}")
# 应用：梯度裁剪（差分隐私）
max_grad_norm = 1.0
clipped_grads = {}
for name, grads in per_sample_grads.items():
# 计算每个样本的梯度范数
grad_norms = grads.flatten(1).norm(2, dim=1)
# 裁剪
clip_coef = max_grad_norm / (grad_norms + 1e-6)
clip_coef = torch.clamp(clip_coef, max=1.0)
# 应用裁剪
clipped_grads[name] = grads * clip_coef.view(-1, *([1] * (grads.dim() - 1)))
print(f"\n{name}梯度裁剪:")
print(f" 裁剪前范数均值: {grad_norms.mean():.4f}")
print(f" 裁剪后范数均值: {(grad_norms * clip_coef).mean():.4f}")
return per_sample_grads, clipped_grads
per_sample_grads, clipped_grads = per_sample_gradients_basic()
# per-sample gradients在差分隐私训练中的应用
def differentially_private_training():
"""差分隐私训练示例"""
class DPTraining:
def __init__(self, model, noise_multiplier=1.0, max_grad_norm=1.0):
self.model = model
self.noise_multiplier = noise_multiplier
self.max_grad_norm = max_grad_norm
def compute_per_sample_grads(self, x_batch, y_batch):
"""计算并裁剪每个样本的梯度"""
def compute_loss(params, x, y):
output = torch.functional.functional_call(self.model, params, x)
return nn.functional.cross_entropy(output, y)
params = dict(self.model.named_parameters())
def compute_and_clip_grad(x, y):
grads = grad(compute_loss)(params, x.unsqueeze(0), y.unsqueeze(0))
# 裁剪梯度
for name, g in grads.items():
grad_norm = g.flatten().norm(2)
clip_coef = self.max_grad_norm / (grad_norm + 1e-6)
clip_coef = min(1.0, clip_coef)
grads[name] = g * clip_coef
return grads
return vmap(compute_and_clip_grad)(x_batch, y_batch)
def add_noise(self, per_sample_grads):
"""添加噪声以实现差分隐私"""
noisy_grads = {}
for name, grads in per_sample_grads.items():
# 计算平均梯度
avg_grad = grads.mean(dim=0)
# 添加高斯噪声
noise = torch.randn_like(avg_grad) * self.noise_multiplier * self.max_grad_norm / len(grads)
noisy_grads[name] = avg_grad + noise
return noisy_grads
def update_model(self, noisy_grads):
"""使用噪声梯度更新模型"""
for name, param in self.model.named_parameters():
if name in noisy_grads:
param.grad = noisy_grads[name]
# 使用优化器更新
return noisy_grads
# 使用示例
model = nn.Linear(20, 10)
dp_trainer = DPTraining(model, noise_multiplier=1.5, max_grad_norm=1.0)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
print("差分隐私训练:")
print(f"Noise multiplier: {dp_trainer.noise_multiplier}")
print(f"Max grad norm: {dp_trainer.max_grad_norm}\n")
for step in range(10):
x = torch.randn(32, 20)
y = torch.randint(0, 10, (32,))
# 计算per-sample梯度并裁剪
per_sample_grads = dp_trainer.compute_per_sample_grads(x, y)
# 添加噪声
noisy_grads = dp_trainer.add_noise(per_sample_grads)
# 更新模型
dp_trainer.update_model(noisy_grads)
optimizer.step()
optimizer.zero_grad()
if step % 3 == 0:
# 计算损失用于显示
output = model(x)
loss = nn.functional.cross_entropy(output, y)
print(f"Step {step}: Loss = {loss.item():.4f}")
print("\n✓ 差分隐私训练完成")
return model, dp_trainer
model, dp_trainer = differentially_private_training()4.6 高阶导数计算
高阶导数在物理信息神经网络、元学习和某些优化算法中非常重要。
python
import torch
import torch.nn as nn
from torch.func import grad, hessian, jacrev
# 二阶导数（海森矩阵）
def second_order_derivatives():
"""二阶导数计算"""
# 方法1: 使用torch.autograd.grad（嵌套）
def compute_hessian_autograd(x):
"""使用autograd计算海森矩阵"""
x = x.clone().requires_grad_(True)
# 定义函数
def f(x):
return (x**3).sum() # f(x) = sum(x_i^3)
# 一阶导数
y = f(x)
first_derivatives = torch.autograd.grad(y, x, create_graph=True)[0]
# 二阶导数（海森矩阵）
hessian = torch.zeros(len(x), len(x))
for i in range(len(x)):
hessian[i] = torch.autograd.grad(
first_derivatives[i], x, retain_graph=True
)[0]
return hessian
# 方法2: 使用torch.func.hessian（推荐）
def compute_hessian_func(x):
"""使用torch.func.hessian计算"""
def f(x):
return (x**3).sum()
return hessian(f)(x)
# 测试
x = torch.tensor([1.0, 2.0, 3.0])
print("海森矩阵计算对比:")
print(f"输入: {x}")
print(f"函数: f(x) = sum(x^3)\n")
hessian_autograd = compute_hessian_autograd(x)
hessian_func = compute_hessian_func(x)
print("方法1 - torch.autograd.grad:")
print(hessian_autograd)
print("\n方法2 - torch.func.hessian:")
print(hessian_func)
print(f"\n结果一致: {torch.allclose(hessian_autograd, hessian_func)}")
# 理论验证
print("\n理论海森矩阵:")
print("对于 f(x) = sum(x_i^3)")
print("H[i,j] = d²f/dx_i dx_j")
print("当 i=j: 6*x_i")
print("当 i≠j: 0")
print(f"所以 H = diag([6, 12, 18])")
return hessian_func
hessian_matrix = second_order_derivatives()
# 高阶导数在物理信息神经网络中的应用
def pinns_higher_order_derivatives():
"""物理信息神经网络中的高阶导数"""
class PhysicsInformedNN(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 50), # (x, t) -> hidden
nn.Tanh(),
nn.Linear(50, 50),
nn.Tanh(),
nn.Linear(50, 1) # -> u(x,t)
)
def forward(self, x, t):
return self.net(torch.cat([x, t], dim=1))
# 一维热传导方程: ∂u/∂t = α ∂²u/∂x²
def heat_equation_residual(model, x, t, alpha=0.1):
"""计算热传导方程的残差"""
x = x.clone().requires_grad_(True)
t = t.clone().requires_grad_(True)
# 预测u(x,t)
u = model(x, t)
# 一阶导数: ∂u/∂t
u_t = torch.autograd.grad(
u, t,
grad_outputs=torch.ones_like(u),
create_graph=True
)[0]
# 一阶导数: ∂u/∂x
u_x = torch.autograd.grad(
u, x,
grad_outputs=torch.ones_like(u),
create_graph=True
)[0]
# 二阶导数: ∂²u/∂x²
u_xx = torch.autograd.grad(
u_x, x,
grad_outputs=torch.ones_like(u_x),
create_graph=True
)[0]
# 残差: ∂u/∂t - α ∂²u/∂x²
residual = u_t - alpha * u_xx
return residual, u, u_t, u_xx
# 使用示例
model = PhysicsInformedNN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
print("物理信息神经网络 - 热传导方程:")
print("方程: ∂u/∂t = 0.1 * ∂²u/∂x²\n")
for step in range(50):
optimizer.zero_grad()
# 生成训练点
x = torch.randn(100, 1)
t = torch.randn(100, 1)
# 计算残差
residual, u, u_t, u_xx = heat_equation_residual(model, x, t)
# 损失: 残差的均方误差
loss = (residual**2).mean()
# 反向传播
loss.backward()
optimizer.step()
if step % 10 == 0:
print(f"Step {step}: Loss = {loss.item():.6f}")
print(f" Mean |∂u/∂t| = {u_t.abs().mean():.4f}")
print(f" Mean |∂²u/∂x²| = {u_xx.abs().mean():.4f}")
print("\n✓ PINN训练完成")
return model
pinn_model = pinns_higher_order_derivatives()
# 三阶及以上导数
def higher_than_second_order():
"""三阶及以上导数计算"""
# 使用torch.func计算高阶导数
def compute_higher_order_derivatives(x):
"""计算高阶导数"""
def f(x):
return torch.sin(x) + x**4
x = torch.tensor(x)
# 一阶导数
f_prime = grad(f)(x)
# 二阶导数
f_double_prime = grad(grad(f))(x)
# 三阶导数
f_triple_prime = grad(grad(grad(f)))(x)
# 四阶导数
f_quadruple_prime = grad(grad(grad(grad(f))))(x)
return f_prime, f_double_prime, f_triple_prime, f_quadruple_prime
x_val = 1.0
derivatives = compute_higher_order_derivatives(x_val)
print("高阶导数计算:")
print(f"函数: f(x) = sin(x) + x^4")
print(f"计算点: x = {x_val}\n")
print("理论导数:")
print(f" f(x) = sin(x) + x^4")
print(f" f'(x) = cos(x) + 4x^3")
print(f" f''(x) = -sin(x) + 12x^2")
print(f" f'''(x) = -cos(x) + 24x")
print(f" f''''(x) = sin(x) + 24\n")
print("计算结果:")
print(f" f'({x_val}) = {derivatives[0]:.6f} (理论: {torch.cos(torch.tensor(x_val)) + 4*x_val**3:.6f})")
print(f" f''({x_val}) = {derivatives[1]:.6f} (理论: {-torch.sin(torch.tensor(x_val)) + 12*x_val**2:.6f})")
print(f" f'''({x_val}) = {derivatives[2]:.6f} (理论: {-torch.cos(torch.tensor(x_val)) + 24*x_val:.6f})")
print(f" f''''({x_val}) = {derivatives[3]:.6f} (理论: {torch.sin(torch.tensor(x_val)) + 24:.6f})")
return derivatives
derivatives = higher_than_second_order()由于篇幅问题，剩下的章节点击这里跳转
第5章 全链路环境标准化部署
5.1 开发环境配置最佳实践
5.2 PyTorch版本兼容性矩阵
5.3 CUDA与cuDNN版本匹配
5.4 分布式训练环境配置
5.5 混合精度训练环境
5.6 容器化部署方案
第6章 实战场景应用案例
6.1 计算机视觉：CNN梯度优化
6.2 自然语言处理：Transformer梯度流
6.3 图神经网络：消息传递梯度
6.4 生成对抗网络：双梯度优化
6.5 强化学习：策略梯度计算
6.6 物理信息神经网络（PINN）
第7章 性能优化与调试技巧
7.1 计算图优化策略
7.2 内存优化技巧
7.3 梯度流可视化
7.4 常见梯度问题诊断
7.5 性能基准测试
第8章 与其他框架对比分析
8.1 PyTorch vs TensorFlow自动微分
8.2 PyTorch vs JAX自动微分
8.3 PyTorch vs MXNet自动微分
8.4 核心差异对比表
8.5 选型指南与适用场景
第9章 未来发展趋势与演进路线
9.1 PyTorch 2.x+新特性展望
9.2 编译优化与torch.compile
9.3 分布式自动微分
9.4 量子机器学习集成
9.5 自动微分标准化
第10章 常见问题与解决方案
第11章 总结与最佳实践
第12章 附录
附录A：完整代码实现示例
附录B：API参考手册
附录C：性能优化checklist
附录D：术语表与参考文献