PyTorch Scala 高校计算机硕士研一课程

高阶梯度计算

PyTorch 的自动求导引擎动态构建计算图并向后遍历以计算梯度，例如 ∂L∂w∂w∂L。这是训练大多数神经网络的基础。然而，某些高级方法需要计算梯度的梯度，这被称为高阶梯度。

考虑一个函数 f(x)f(x)。它的一阶导数是 f′(x)=dfdxf′(x)=dxdf。二阶导数是 f′′(x)=d2fdx2f′′(x)=d**x2d2f，它简单来说就是一阶导数的导数。同样地，我们可以计算三阶、四阶及更高阶导数。在多变量函数（例如带有参数 θθ 的神经网络损失函数 L(θ)L(θ)）的上下文中，我们通常处理偏导数。一阶梯度构成梯度向量 ∇L∇L。高阶导数涉及 Hessian 矩阵（二阶偏导数矩阵，∇2L∇2L）或更高阶张量等结构。

PyTorch 的自动求导引擎能够处理这些计算。虽然标准的 .backward() 方法主要用于一阶梯度，但 torch.autograd.grad 函数式接口提供了高阶微分所需的灵活性。

为何计算高阶梯度？

计算高阶梯度对一些高级应用非常重要：

1. 优化算法： 像牛顿法或信任区域法这样的方法使用二阶信息（Hessian 矩阵）来可能实现比 SGD 或 Adam 等一阶方法更快的收敛。虽然对于大型网络来说，计算完整的 Hessian 矩阵通常不可行，但通过高阶自动微分可以高效地计算 Hessian-向量积 (∇2Lv∇2Lv)，并被用于一些优化策略中。
2. 曲率分析： 二阶导数（Hessian 矩阵）描述了损失函数的曲率。分析这种曲率可以提供关于优化过程、泛化特性以及局部最小值或鞍点存在情况的理解。
3. 元学习： 像模型无关元学习 (MAML) 这样的算法涉及根据模型在特定任务上经过一次或多次梯度更新后的表现来优化其参数。这需要对梯度更新步骤本身进行微分，因此需要计算梯度的梯度。
4. 正则化技术： 某些正则化项明确依赖于梯度范数或二阶导数。例如，带梯度惩罚的 Wasserstein GAN (WGAN-GP) 中的梯度惩罚需要计算判别器输出相对于其输入的梯度的范数。
5. 物理信息神经网络 (PINNs)： PINNs 将物理定律（常以偏微分方程 (PDEs) 形式表示）融入损失函数。这些 PDE 常涉及网络输出相对于其输入坐标（例如，时间和空间）的二阶或更高阶导数。

使用 torch.autograd.grad 计算高阶梯度

在 PyTorch 中计算高阶梯度的主要工具是 torch.autograd.grad。与 tensor.backward() 方法隐式计算所有需要梯度的叶节点梯度不同，torch.autograd.grad 更显式。

其基本签名如下：

import torch

// 计算高阶梯度的基本签名
torch.autograd.grad(
    outputs,        // 要进行微分的标量或张量
    inputs,         // 计算梯度时所依据的张量
    grad_outputs=None, // 损失函数对 'outputs' 的梯度（用于向量-雅可比积）
    retain_graph=None, // 如果为 True，保留图；否则释放。
    create_graph=false, // 如果为 True，为梯度计算本身构建
    allow_unused=false
)

对于高阶梯度，重要的参数是 create_graph=True。当您使用 torch.autograd.grad 并设置 create_graph=True 来计算一阶梯度时，PyTorch 不仅计算梯度，还会构建必要的图结构，以便您稍后可以对这次梯度计算进行再次微分。如果 create_graph=False（默认值），梯度计算被视为一个终端操作；生成的梯度只是张量，没有将它们通过微分过程连接回原始参数的任何历史记录。

我们来看一个简单的例子。假设我们有 y=x3y=x3。我们想计算 dydx=3x2dxd**y=3x2 和 d2ydx2=6xd**x2d2y=6x。

import torch

// 输入张量需要梯度
val x = torch.tensor([2.0], requires_grad=true)

// 第一次计算: y = x^3
val y = x**3
println(f"y = {y.item()}")

// 计算一阶导数: dy/dx
// 使用 create_graph=True 以允许计算高阶梯度
val grad_y_x = torch.autograd.grad(outputs=y, inputs=x, create_graph=true)[0]
println(f"x={x.item()} 处的 dy/dx: {grad_y_x.item()}") // 应该是 3 * (2^2) = 12

// grad_y_x 现在是一个带有自身计算图的张量
println(f"梯度张量 requires_grad: {grad_y_x.requires_grad}")

// 计算二阶导数: d^2y/dx^2 = d/dx (dy/dx)
// 我们对*一阶梯度* (grad_y_x) 相对于 x 进行微分
// 除非我们想要三阶梯度，否则这里不需要 create_graph=True
val grad2_y_x2 = torch.autograd.grad(outputs=grad_y_x, inputs=x)[0]
println(f"x={x.item()} 处的 d^2y/dx^2: {grad2_y_x2.item()}") // 应该是 6 * 2 = 12

// 检查二阶导数的 requires_grad 状态
println(f"二阶导数张量 requires_grad: {grad2_y_x2.requires_grad}")

请注意，grad_y_x 的 requires_grad=True，因为我们在其计算过程中指定了 create_graph=True。这允许我们再次以 grad_y_x 作为输出调用 torch.autograd.grad。最终的 grad2_y_x2 的 requires_grad=False，因为我们在第二次调用中未指定 create_graph=True。

图修改示意

当使用 create_graph=True 时，反向传播过程本身会将节点添加到计算图中。

考虑 y=x2y=x2，所以 dydx=2xdxd**y=2x。

1. 前向传播： x -> pow(2) -> y
2. 反向传播 (create_graph=False)： 计算梯度 (2x2x) 并将其作为与用于计算它的图分离的新张量返回。
3. 反向传播 (create_graph=True)： 计算梯度 (2x2x)，但会将表示该梯度如何计算的操作添加到图中：x -> pow(2) -> y；grad_y -> MulBackward (使用保存的 x) -> grad_x。输出 grad_x 被连接到这个扩展图上。

该图对比了 torch.autograd.grad 在 create_graph=False (中间) 和 create_graph=True (右侧) 时的结果。当 create_graph=True 时，计算出的梯度 grad_x 通过梯度计算操作 (PowBackward) 保持与图的连接，从而允许进一步微分。

示例：Hessian-向量积

我们来计算一个简单函数 f(w1,w2)=w12sin⁡(w2)f(w1,w2)=w12sin(w2) 的 Hessian-向量积 (HVP)。梯度为 ∇f=[∂f∂w1,∂f∂w2]=[2w1sin⁡(w2),w12cos⁡(w2)]∇f=[∂w1∂f,∂w2∂f]=[2w1sin(w2),w12cos(w2)]。Hessian 矩阵为 ∇2f=(∂2f∂w12∂2f∂w1∂w2∂2f∂w2∂w1∂2f∂w22)=(2sin⁡(w2)2w1cos⁡(w2)2w1cos⁡(w2)−w12sin⁡(w2))∇2f=(∂w12∂2f∂w2∂w1∂2f∂w1∂w2∂2f∂w22∂2f)=(2sin(w2)2w1cos(w2)2w1cos(w2)−w12sin(w2))。

我们希望计算 (∇2f)v(∇2f)v，其中 vv 为某个向量，且无需显式构造 ∇2f∇2f。通过两次 torch.autograd.grad 调用即可实现。主要思路是 (∇2f)v=∇(∇f⋅v)(∇2f)v=∇(∇f⋅v)，其中 ∇f⋅v∇f⋅v 是点积（一个标量）。

import torch

val w = torch.tensor(Seq(1.0, torch.pi / 2.0), requires_grad=true) // w1=1，w2=pi/2
val v = torch.tensor(Seq(0.5, 1.0)) // 一个任意向量

// 定义函数
val f = w[0]**2 * torch.sin(w[1])

// 计算一阶梯度: grad_f = nabla f
val grad_f = torch.autograd.grad(f, w, create_graph=true)[0]
// 预期 grad_f: [2*1*sin(pi/2), 1^2*cos(pi/2)] = [2, 0]
println(f"梯度 (nabla f): {grad_f}")

// 计算点积: grad_f_dot_v = (nabla f) . v
// 这个操作需要成为图的一部分，以便进行第二次微分
val grad_f_dot_v = torch.dot(grad_f, v)
println(f"点积 (nabla f . v): {grad_f_dot_v}") // 预期: 2*0.5 + 0*1 = 1.0

// 计算点积相对于 w 的梯度: nabla (nabla f . v)
// 这得到 Hessian-向量积 (nabla^2 f) v
val hvp = torch.autograd.grad(grad_f_dot_v, w)[0]
// 预期 Hessian: [[2*sin(pi/2), 2*1*cos(pi/2)], [2*1*cos(pi/2), -1^2*sin(pi/2)]]
// = [[2, 0], [0, -1]]
// 预期 HVP: [[2, 0], [0, -1]] @ [0.5, 1.0] = [2*0.5 + 0*1, 0*0.5 + (-1)*1] = [1.0, -1.0]
println(f"Hessian-向量积 (nabla^2 f) v: {hvp}")

这种方法避免了显式生成可能非常大的 Hessian 矩阵，仅需要向量积和梯度计算，对于大型模型来说内存效率更高。

注意事项

• 计算成本： 计算高阶梯度的开销比一阶梯度大。每次调用 torch.autograd.grad 并设置 create_graph=True，在后续反向传播中，图的遍历深度实质上会翻倍。
• 内存占用： 存储高阶导数所需的图会消耗更多内存。
• 二次反向传播： 计算二阶导数的过程有时被称为“二次反向传播”。

了解如何使用 torch.autograd.grad 和 create_graph=True 标志计算高阶梯度，实现了在优化、模型分析以及 PyTorch 框架内实现复杂算法（如元学习和物理信息建模）方面的一系列高级能力。

梯度查看与图可视化

了解梯度如何在网络中传递，对于调试和优化必不可少。当模型表现异常或训练停滞时，检查梯度和其对应的计算图通常能提供有用的线索。PyTorch在对需要梯度的张量执行操作时，会动态地构建这个计算图。这里将介绍查看这些梯度和可视化图结构的方法。

获取与分析梯度

在调用loss.backward()之后，PyTorch会计算损失相对于计算图中所有requires_grad=True且参与了损失计算的张量的梯度。这些梯度会累积在对应的叶子张量（通常是模型参数或输入）的.grad属性中。

import torch

// 示例设置
val w = torch.randn(5, 3, requires_grad=true)
val x = torch.randn(3, 2)
val y_true = torch.randn(5, 2)

// 前向传播
val y_pred = w @ x
val loss = torch.nn.functional.mse_loss(y_pred, y_true)

// 反向传播
loss.backward()

// 查看w中累积的梯度
println("Gradient for w:\n", w.grad)

// 非叶子张量或requires_grad=False的张量的梯度通常为None
println("Gradient for x:", x.grad) // 输出: None (默认requires_grad=False)
println("Gradient for y_pred:", y_pred.grad) // 输出: None (非叶子张量，默认不保留梯度)

检查.grad时常见的几种情况：

1. 梯度为None： 如果张量的.grad在调用.backward()后为None，通常表示：
   • 该张量未设置requires_grad=True。
   • 该张量未参与到导致损失的计算图中（例如，它是在with torch.no_grad():块中创建或使用.detach()进行分离的）。
   • 它是非叶子张量，PyTorch默认不保存中间梯度以节省内存。如果需要查看中间结果的梯度，请使用tensor.retain_grad()。
2. 梯度消失： 梯度变得非常小（例如，10−810−8或更小），通常接近零。这导致图后面（更接近输入端）的权重无法有效更新。在深度网络或使用sigmoid等激活函数的网络中很常见，特别是在没有批量归一化或残差连接等方法的情况下。
3. 梯度爆炸： 梯度变得过大（例如，108108或NaN）。这导致训练不稳定、权重更新幅度大，并且损失或权重中常出现NaN值。梯度裁剪（第三章介绍）是一种常见的缓解策略。

你可以编程方式检查这些问题：

// 检查None梯度（假设'model'是你的torch.nn.Module实例）
for (name, param) <- model.named_parameters() {
    if param.grad == None {
        println(f"Parameter {name} has no gradient.")
    }
}

// 检查梯度消失/爆炸
var max_grad_norm = 0.0
var min_grad_norm = Float.PositiveInfinity
var nan_detected = false
for param in model.parameters():
    if param.grad != None:
        val grad_norm = param.grad.norm().item()
        if torch.isnan(param.grad).any():
             nan_detected = true
             println(f"NaN gradient detected in parameter: {param.size()}") // 可能需要更具体的识别
        max_grad_norm = math.max(max_grad_norm, grad_norm)
        min_grad_norm = math.min(min_grad_norm, grad_norm)

println(f"Max gradient norm: {max_grad_norm:.4e}")
println(f"Min gradient norm: {min_grad_norm:.4e}")
if nan_detected then
    println("Warning: NaN gradients detected!") // 警告：检测到NaN梯度！

使用Hook进行细粒度检查

为在反向传播期间进行更详细的分析，PyTorch提供了hook（钩子）。Hook是可以在特定事件发生时（例如张量梯度计算或模块的前向/反向传播）注册执行的函数。

张量Hook (register_hook)

你可以直接在张量上注册一个hook。当该特定张量的梯度被计算时，这个hook函数将执行。hook函数将梯度作为其唯一参数接收。

import torch.*

def print_grad_hook(grad: torch.Tensor):
    println(f"Gradient received: shape={grad.shape}, norm={grad.norm():.4f}")

val x = torch.randn(3, 3, requires_grad=true)
val y = x.pow(2).sum()

// 在张量x上注册hook
val hook_handle = x.register_hook(print_grad_hook)

// 计算梯度
y.backward()

// hook函数（print_grad_hook）会自动调用
// 输出将包含类似以下内容：
// Gradient received: shape=torch.Size([3, 3]), norm=9.5930

// 不再需要时应移除hook以避免内存泄漏
hook_handle.remove()

// 你也可以在hook中修改梯度，但请谨慎使用：
def scale_grad_hook(grad: torch.Tensor):
  // 示例：将梯度减半
  return grad * 0.5

// x.register_hook(scale_grad_hook)
// y.backward() // 现在存储在x.grad中的梯度将减半

Hook对于调试网络的特定部分非常有用。你可以记录梯度统计信息，在NaN值出现时精准检查它们，甚至实时修改梯度（尽管修改梯度通常较不常见且需要仔细斟酌）。

模块Hook

你也可以在torch.nn.Module实例上注册hook，以便在前向传播期间检查输入和输出，或在反向传播期间检查梯度。

• register_forward_pre_hook(hook): 在模块的forward方法之前执行。接收参数(module, input)。
• register_forward_hook(hook): 在模块的forward方法之后执行。接收参数(module, input, output)。
• register_full_backward_hook(hook): 在为模块的输入和输出计算完梯度后执行。接收参数(module, grad_input, grad_output)。grad_input是一个包含模块输入梯度的元组，grad_output是一个包含模块输出梯度的元组。

import torch
import torch.nn as nn

class SimpleNet extends nn.Module:
    def __init__(self):
        super().__init__()
        val linear1 = nn.Linear(10, 5)
        val relu = nn.ReLU()
        val linear2 = nn.Linear(5, 1)

    def forward(x: torch.Tensor):
        x = linear1(x)
        x = relu(x)
        x = linear2(x)
        return x

val model = SimpleNet()
val input_tensor = torch.randn(4, 10, requires_grad=true)

def backward_hook(module: nn.Module, grad_input: Seq[torch.Tensor], grad_output: Seq[torch.Tensor]):
    println(f"\nModule: {module.__class__.__name__}")
    println("  grad_input shapes: ", [g.shape if g is not None else None for g in grad_input])
    println("  grad_output shapes:", [g.shape if g is not None else None for g in grad_output])

// 在linear2层上注册hook
val hook_handle_bwd = model.linear2.register_full_backward_hook(backward_hook)

// 前向和反向传播
val output = model(input_tensor)
val target = torch.randn(4, 1)
val loss = nn.functional.mse_loss(output, target)
loss.backward()

// 输出将显示流经linear2的反向梯度形状
// Module: Linear
//   grad_input shapes:  [torch.Size([4, 5]), torch.Size([5]), None] （输入、权重、偏置）如果bias=False，偏置梯度可能为None
//   grad_output shapes: [torch.Size([4, 1])]

// 清理hook
hook_handle_bwd.remove()

模块hook对于理解梯度如何逐层传播，或诊断大型网络中特定模块的问题特别有用。

可视化计算图

尽管hook能让你以数值方式检查梯度，但可视化计算图可以提供结构性概览。这有助于理解操作与参数之间的依赖关系，确认你的模型架构，或发现意外连接。

使用torchviz

一个用于基础图可视化的流行第三方库是torchviz。它使用graphviz库来渲染在反向传播期间生成的图。

你通常会在输出张量（通常是损失）上调用torchviz.make_dot，以可视化其梯度计算图。它返回一个graphviz.Digraph对象。

// 要求：pip install torchviz graphviz
import torch
import torchviz

// 简单示例
val a = torch.tensor([2.0], requires_grad=true)
val b = torch.tensor([3.0], requires_grad=true)
val c = a * b
val d = c + a
val L = d.mean() // 最终标量输出

// 生成图可视化对象
// params可用于突出显示特定参数
val graph = torchviz.make_dot(L, params=Map("a" -> a, "b" -> b))

// 要查看图，你可以将其渲染到文件或在Jupyter等环境中显示
// graph.render("computation_graph", format="png") // 保存为computation_graph.png
// display(graph) // 在Jupyter环境中

// 为演示目的，我们打印Graphviz源代码
// print(graph.source)

a参数MulBackward0grad_fnAddBackward0grad_fnb参数grad_fnAccumulateGradAccumulateGradgrad_fnabAccumulateGradMeanBackward0grad_fnAccumulateGradd

一个由torchviz生成的简单计算图。椭圆代表张量（参数高亮显示），方框代表反向操作（grad_fn）。箭头表示反向传播期间的梯度流向。

torchviz提供了反向图的清晰高层视图，非常适合理解依赖关系和梯度计算流程。

使用TensorBoard

PyTorch内置支持TensorBoard，这是一个来自TensorFlow的强大可视化工具包。你可以使用torch.utils.tensorboard.SummaryWriter记录计算图（以及许多其他内容，如标量、图像、直方图）。

import torch
import torch.nn as nn
import torch.utils.tensorboard.SummaryWriter

// 再次定义一个简单模型
class SimpleNet extends nn.Module:
    def __init__(self):
        super().__init__()
        val layer1 = nn.Linear(5, 3)
        val relu = nn.ReLU()
        val layer2 = nn.Linear(3, 1)
    def forward(x: torch.Tensor):
        return layer2(relu(layer1(x)))

val model = SimpleNet()
val dummy_input = torch.randn(1, 5) // 提供一个示例输入

// 创建一个SummaryWriter实例（默认日志保存到./runs/）
val writer = new SummaryWriter("runs/graph_demo")

// 将图添加到TensorBoard
// writer需要模型和一个示例输入张量
writer.add_graph(model, dummy_input)
writer.close()

// 要查看图：
// 1. 确保已安装tensorboard (pip install tensorboard)
// 2. 在你的终端运行`tensorboard --logdir=runs/graph_demo`
// 3. 在浏览器中打开提供的URL（通常是http://localhost:6006/）
// 4. 导航到“Graphs”选项卡。

TensorBoard直接在你的浏览器中提供了一个交互式图可视化环境。它通常显示一个更详细的图，包括模块范围、参数节点和操作节点。虽然对于非常大的模型可能显得过于复杂，但其交互性允许你展开和折叠图的部分，使其比静态图像更容易浏览复杂的架构。

实际考量

• 开销： 注册许多hook，特别是执行复杂计算或大量I/O（如打印）的hook，会显著减慢训练速度。主要将其用于调试，除非必要，否则不要在生产训练循环中使用。
• 图的复杂性： 对于非常深或复杂的模型，完整的计算图会变得非常庞大且难以视觉解读。将可视化工作集中在与调试任务相关的特定模块或子图上。
• 动态图： 请记住PyTorch的图是动态的。可视化的图对应于使用给定输入执行的特定前向传播。如果你的模型具有数据依赖的控制流（例如，影响所用层的if语句），图在不同迭代或不同输入之间可能发生变化。

有效检查梯度和可视化计算图是PyTorch高级开发中不可或缺的技能。它们能让你对框架有更深刻的理解，实现有针对性的调试，并做出明智的优化决策。下一章将在此基础上，实现复杂的网络架构。

内存管理考量

随着您构建和训练日益复杂的模型，管理内存成为开发和调试的一个重要方面，尤其是在使用GPU时，因为GPU的内存通常比CPU更有限。了解PyTorch如何处理内存分配以及您的操作如何影响内存，对于提高效率和避免常见的内存不足错误很必要。PyTorch的内存管理机制，以及它们与张量结构和自动求导过程的相互作用，在此进行审视。

张量存储与内存布局

其核心是，PyTorch张量（torch.Tensor）是对由torch.Storage对象管理的连续内存块的视图。Storage对象保存实际的数值数据，而Tensor对象包含形状（大小）、步长和数据类型（dtype）等元数据，以及它在Storage中的位置信息。

多个张量可以共享同一个底层Storage。例如，对张量进行切片或使用view()等操作通常会创建一个新的张量对象，它指向相同的存储但具有不同的元数据。

import torch.*

// 创建一个张量；PyTorch分配存储空间
val x = torch.randn(2, 3)
println(f"x storage: {x.storage().data_ptr()}")

// 切片操作会创建一个共享存储的新张量视图
val y = x[0, :]
println(f"y storage: {y.storage().data_ptr()}") // 相同的指针
println(f"Do x and y share storage? {x.storage().data_ptr() == y.storage().data_ptr()}")

// 修改y会影响x，因为它们共享存储
y.fill_(1.0)
println("修改y后x的值：\n", x)

这种存储共享非常高效，因为它避免了不必要的数据复制。然而，了解这一点很重要，尤其是在执行原地操作时。

张量在内存中的布局由其步长决定。如果张量的元素在内存中逐行（对于二维张量）顺序排列且没有间隙，则认为该张量是连续的。非连续张量可能由转置或某些类型的索引操作产生。

// 连续张量
val a = torch.arange(6).reshape(2, 3)
println(f"a is contiguous: {a.is_contiguous()}, Stride: {a.stride()}") // 步长: (3, 1)

// 转置会创建非连续视图
val b = a.t()
println(f"b is contiguous: {b.is_contiguous()}, Stride: {b.stride()}") // 步长: (1, 3)

// 访问元素仍然正确，但内存访问模式不同
println("b:\n", b)

// 某些PyTorch函数需要连续张量
// 尝试对非连续张量进行view等操作可能会失败
try:
    b.view(-1)
catch RuntimeError(e) {
    println(f"\nError viewing non-contiguous tensor: {e}")
}

// 使用 .contiguous() 获取连续副本
val c = b.contiguous()
println(f"c is contiguous: {c.is_contiguous()}, Stride: {c.stride()}") // 步长: (2, 1)
println("c (contiguous version of b):\n", c)
println(f"Does b and c share storage? {b.storage().data_ptr() == c.storage().data_ptr()}") // 否，新的存储空间

虽然PyTorch操作通常能正确处理非连续张量，但某些底层操作或接口（例如导出到NumPy或某些自定义扩展）可能需要连续数据。如果原始张量不是连续的，调用.contiguous()会创建一个带有全新、连续数据副本的新张量。这会产生内存复制开销。

数据类型（dtype）也直接影响内存使用。一个torch.float32张量每个元素使用4字节，而torch.float16使用2字节，torch.int64使用8字节。选择合适的数据类型是提高内存效率的基本要求。

PyTorch缓存内存分配器

使用CUDA API（cudaMalloc、cudaFree）在GPU上分配和释放内存可能很慢。为了缓解这个问题，PyTorch为GPU张量采用了一个缓存内存分配器。当一个张量被释放时（例如，超出作用域且其引用计数降至零），它所占用的内存不一定立即返回给GPU操作系统。相反，PyTorch将此内存块保留在缓存中。

当需要分配新张量时，PyTorch首先检查其缓存中是否有大小合适的空闲块。如果找到，它会重用该块，避免了对CUDA驱动程序的昂贵调用。这显著加快了张量的创建和删除速度，而这在训练期间经常发生。

CUDA 驱动程序PyTorch 分配器分配/释放活跃张量提供内存缓存块 (非活跃)保留已释放内存释放内存重用内存

PyTorch缓存分配器与CUDA驱动程序和张量内存交互的简化视图。

您可以查看缓存分配器的状态：

• torch.cuda.memory_allocated()：返回默认设备上张量当前占用的总GPU内存（以字节为单位）。
• torch.cuda.memory_reserved() 或 torch.cuda.memory_cached()（已弃用）：返回缓存分配器管理的总GPU内存（包括已分配的张量和缓存的空闲块）。
• torch.cuda.max_memory_allocated()：返回从开始执行或上次重置以来，在任何时间点张量占用的最大GPU内存。
• torch.cuda.reset_peak_memory_stats()：重置峰值内存计数器。
• torch.cuda.memory_summary()：提供已分配和缓存内存的详细报告，通常有助于发现碎片问题。

有时，您可能希望清除缓存内存，也许是为了使其可供其他GPU应用程序或库使用。您可以使用torch.cuda.empty_cache()。

// 需要GPU
if torch.cuda.is_available() then
    val device = torch.device("cuda")
    println(f"Initial allocated: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MiB")
    println(f"Initial reserved:  {torch.cuda.memory_reserved(device) / 1024**2:.2f} MiB")

    // 分配一些张量
    val t1 = torch.randn(1024, 1024, device=device)
    val t2 = torch.randn(512, 512, device=device)
    println(f"\nAfter allocation:")
    println(f"Allocated: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MiB")
    println(f"Reserved:  {torch.cuda.memory_reserved(device) / 1024**2:.2f} MiB")

    // 删除张量
    t1.delete()
    t2.delete()
    println(f"\nAfter deleting tensors (before empty_cache):")
    // 已分配内存减少，但由于缓存，保留内存仍然很高
    println(f"Allocated: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MiB")
    println(f"Reserved:  {torch.cuda.memory_reserved(device) / 1024**2:.2f} MiB")

    // 清除缓存
    torch.cuda.empty_cache()
    println(f"\nAfter empty_cache:")
    // 保留内存也减少（尽管可能由于内部分配而不降为零）
    println(f"Allocated: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MiB")
    println(f"Reserved:  {torch.cuda.memory_reserved(device) / 1024**2:.2f} MiB")
else 
    println("CUDA不可用，跳过GPU内存示例。")

重要提示： torch.cuda.empty_cache() 不会释放当前被活跃张量使用的内存。它只释放未被任何张量引用的缓存块。它主要用于将内存释放回系统以供其他进程使用，而不是在张量仍然存在的情况下减少您正在运行的PyTorch脚本的内存占用。它还会产生性能开销，因为后续分配将需要再次请求驱动程序。

缓存分配器的一个副作用是碎片化。如果您分配和释放不同大小的张量，缓存最终可能持有许多小的、非连续的空闲块。即使这些缓存块的总大小很大，您也可能无法分配一个大的连续块，从而导致内存不足（OOM）错误。torch.cuda.memory_summary()可以帮助诊断碎片问题。

自动求导和内存

自动求导引擎显著影响内存使用。为了在反向传播期间计算梯度，自动求导通常需要存储作为计算图一部分的中间激活值（前向操作的输出）。

• 计算图： 当对需要梯度的张量（requires_grad=True）执行操作时，PyTorch会构建一个图，存储这些操作以及对所涉及张量的引用。这些引用会使张量在内存中保持活跃，即使它们在您的Python代码中可能看起来已超出作用域。
• 反向传播： 在loss.backward()期间，自动求导会反向遍历此图。它使用存储的中间值来计算梯度。一旦梯度被计算并且在反向传播中不再需要进行进一步计算时，持有相应中间激活值的缓冲区通常会被释放。
• retain_graph=True： 如果您调用backward(retain_graph=True)，即使在反向传播完成后，PyTorch也会保留图和中间激活缓冲区。这允许您多次调用backward()（例如，计算不同损失相对于相同参数的梯度），但这代价是占用可能大量的内存。仅在必要时使用它。
• torch.no_grad()： 将代码包裹在with torch.no_grad():块中会向PyTorch发出信号，表明此块内的操作不应被自动求导跟踪。这可以防止为这些操作创建计算图，并避免存储中间激活值，从而节省大量内存。在验证或推理循环中使用此上下文管理器是标准做法。
• .detach()： 对张量调用.detach()会创建一个新张量，它共享相同的存储空间但与计算图分离。它不需要梯度，并且不涉及它的操作将不会被跟踪。如果您需要使用张量的值而不跟踪其历史记录（例如，用于日志记录或绘图），这很有用。

考虑这个简单示例：

// 设置
val a = torch.randn(100, 100, requires_grad=true)
val b = torch.randn(100, 100, requires_grad=true)

// 被自动求导跟踪的操作
val c = a * b
val d = c.sin()
val loss = d.mean()

// 中间张量'c'和'd'被保留在内存中
// 因为反向传播需要它们。

// 调用backward会释放缓冲区（除非retain_graph=True）
loss.backward() // 计算a和b的梯度

// 现在，让我们尝试不跟踪梯度
with torch.no_grad():
    val c_no_grad = a * b // 操作已执行，但未被跟踪
    val d_no_grad = c_no_grad.sin()
    val loss_no_grad = d_no_grad.mean()
    // PyTorch不需要为未来的反向传播存储'c_no_grad'
    // 中间结果的内存可能更早被释放。

println(f"a的梯度：{'存在' if a.grad is not None else '无'}")
// loss_no_grad.backward() // 这将引发错误，因为历史记录未被跟踪。

高效内存使用的方法

以下是实用的方法：

1. 作用域和del： 当对象不再被引用时，Python的垃圾回收器会回收内存。确保不再需要的大张量超出作用域。如果需要，可以使用del语句明确删除引用，尤其是在可能进行内存密集型操作（如backward()）或分配新的大张量之前。

def process_data(data): intermediate = data * 2 # 大型中间张量 result = intermediate.sum() # 如果不删除，'intermediate’可能会在内存中停留更长时间 del intermediate # 明确删除引用 return result ```

1. 原地操作： 以单个下划线（_）结尾的操作，如add_()、relu_()，会直接修改张量，而不是创建新张量。这样可以通过避免为结果分配新张量来节省内存。 注意： 原地修改计算梯度所需的张量可能会破坏反向传播。自动求导会跟踪原地操作，如果检测到此类修改干扰梯度计算，就会引发错误。请谨慎使用它们，通常在图中是叶子节点或您确定不会影响所需梯度的张量上使用。

val x = torch.randn(1000, 1000) y = torch.randn(1000, 1000)

非原地操作：创建一个新张量z

val z = x + y

原地操作：直接修改x，为结果张量节省内存

x.add_(y) # x现在包含x + y的结果 ```

1. 梯度检查点（激活检查点）： 对于具有非常深层结构的模型，如果存储所有中间激活值会消耗过多内存，梯度检查点提供了一种权衡。它在前向传播期间只存储一部分激活值，而不是全部。在反向传播期间，它会即时重新计算必要的激活值。这会使用更多的计算时间，但显著减少峰值内存使用。PyTorch为此提供了torch.utils.checkpoint.checkpoint。
2. 混合精度训练： 使用torch.float16或torch.bfloat16等低精度数据类型，与torch.float32相比，存储激活值、梯度和参数所需的内存减少一半。torch.cuda.amp（自动混合精度）等库有助于有效管理这一点（第3章介绍）。
3. 数据加载和批大小： 确保您的数据加载流程高效。如果遇到OOM错误，减小批大小通常是第一步，因为激活值及其梯度会随批大小线性增长。

内存问题调试

• 内存不足（OOM）错误： 当您遇到CUDA OOM错误时，错误消息本身通常会说明请求了多少内存以及有多少可用内存。
  • 使用torch.cuda.memory_summary()来查看已分配块和缓存碎片的分布。即使总空闲内存看起来足够，高度碎片化也可能导致OOM。
  • 系统地减小批大小。
  • 检查模型大小和复杂度。
  • 在训练循环的不同位置插入torch.cuda.memory_allocated()的打印语句，以找出内存使用量激增的地方。
  • 使用PyTorch分析器（第4章介绍）获取每个操作符的内存使用详细分类。
• 内存泄漏： 如果内存使用量在训练迭代过程中持续增长而不稳定，您可能存在内存泄漏。这通常发生在带有计算历史的张量在torch.no_grad()上下文之外被无意中累积到列表或字典中时。
  • 泄漏示例：在不分离的情况下将损失存储在列表中：all_losses.append(loss)而不是all_losses.append(loss.item())或all_losses.append(loss.detach())。存储原始的loss张量会使其整个计算图保持活跃。
  • 仔细检查张量在迭代之间如何存储。如果您需要张量值而不跟踪其历史记录，可以使用.item()从单元素张量获取Python数字，或者使用.detach()。

有效的内存管理通常是一个迭代的过程，它需要理解模型行为，应用适当的方法，并使用PyTorch的工具检查和调试内存使用。扎实掌握这些知识在扩展到更大的数据集和更复杂的架构时是不可或缺的。

实践操作：构建自定义自动求导函数

虽然PyTorch的自动微分功能可以处理大多数标准操作，但您会遇到需要自定义梯度逻辑的情况。这可能是因为您正在实现一项新颖的操作，优化特定计算，或者处理自动求导无法直接推导出梯度的函数。获得使用torch.autograd.Function定义自己可微分操作的实践经验。

构建模块：torch.autograd.Function

定义具有特定梯度规则的自定义操作的核心机制是继承torch.autograd.Function。该类要求您实现两个静态方法：

1. forward(): 此方法执行操作的实际计算。它接收输入张量，并可以接受额外参数。重要的是，它还接收一个上下文对象ctx，该对象作为通向backward方法的桥梁。您可以使用ctx.save_for_backward()来存储后续梯度计算所需的任何张量。它应返回操作的输出张量。
2. backward(): 此方法定义梯度计算。它接收上下文对象ctx（包含从forward保存的张量）以及损失相对于forward方法输出（grad_output）的梯度。其职责是计算并返回损失相对于forward方法每个输入的梯度。返回的梯度数量和顺序必须与forward的输入数量和顺序匹配。如果一个输入不需要梯度（例如，它不是张量或requires_grad=False），您应为其对应的梯度返回None。

示例：实现一个截断ReLU函数

我们来实作一个自定义激活函数：截断ReLU。该函数行为类似标准ReLU，但将最大输出值限制在特定阈值。

从数学上说，对于截断值CC：

截断ReLU(x,C)=min⁡(max⁡(0,x),C)截断ReLU(x,C)=min(max(0,x),C)

相对于xx的导数是：

∂∂x截断ReLU(x,C)={1如果 0<x<C0否则∂x∂截断ReLU(x,C)={10如果 0<x<C否则

现在，我们使用torch.autograd.Function来实作它。

import torch

class ClippedReLUFunction(torch.autograd.Function):
    """
    实现截断ReLU函数：min(max(0, x), clip_val)。
    """

    @staticmethod
    def forward(ctx, input_tensor, clip_val):
        """
        前向传播：计算截断ReLU。

        参数：
            ctx: 用于保存信息供反向传播的上下文对象。
            input_tensor: 输入张量。
            clip_val: 输出的截断最大值。

        返回：
            应用截断ReLU后的输出张量。
        """
        // 确保clip_val为浮点数以便一致比较
        val clip_val = float(clip_val)
        // 保存输入张量和clip_val供反向传播使用
        // 我们只需要输入张量来计算梯度掩码
        ctx.save_for_backward(input_tensor)
        // 将非张量参数直接存储在ctx上
        ctx.clip_val = clip_val

        // 应用截断ReLU操作
        output = input_tensor.clamp(min=0, max=clip_val)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        反向传播：计算截断ReLU的梯度。

        参数：
            ctx: 带有保存信息的上下文对象。
            grad_output: 损失相对于此函数输出的梯度。

        返回：
            相对于input_tensor的梯度，相对于clip_val的梯度（无）
        """
        // 检索已保存的张量和值
        val input_tensor, = ctx.saved_tensors
        val clip_val = ctx.clip_val

        // 根据输入值范围创建梯度掩码
        // 当 0 < 输入 < clip_val 时梯度为1，否则为0
        val grad_input_mask = (input_tensor > 0) & (input_tensor < clip_val)
        val grad_input = grad_output * grad_input_mask.float()

        // 由于clip_val是一个超参数，因此不需要计算其梯度，
        // 它不是我们通常进行微分的输入张量。
        // 对于非张量输入或不需要梯度的输入，
        // 返回None作为其梯度。
        return grad_input, None

// 辅助函数，使其更易于像标准PyTorch函数一样使用
def clipped_relu(input_tensor, clip_val=1.0):
    """逐元素应用截断ReLU函数。"""
    return ClippedReLUFunction.apply(input_tensor, clip_val)

// 使用示例
val x = torch.randn(5, requires_grad=true, dtype=torch.float64) // 使用float64以获得gradcheck所需更高精度
val clip_value = 2.0
val y = clipped_relu(x, clip_value)
val z = y.mean() // 下游计算示例

// 计算梯度
z.backward()

println("输入张量 (x):\n", x)
println("截断输出 (y):\n", y)
println("平均输出 (z):\n", z)
println("x的梯度 (x.grad):\n", x.grad)

在此代码中：

• ClippedReLUFunction继承自torch.autograd.Function。
• forward计算 y=min⁡(max⁡(0,x),C)y=min(max(0,x),C)，使用ctx.save_for_backward(input_tensor)保存梯度计算所需的输入张量x，并将非张量clip_val直接保存到ctx上。
• backward使用ctx.saved_tensors检索input_tensor。它计算梯度掩码（如果0<x<C0<x<C则为11，否则为00），并将其与传入梯度grad_output逐元素相乘。它返回input_tensor的计算梯度，并为clip_val返回None，因为clip_val不是需要梯度的张量输入。
• clipped_relu辅助函数提供了一个用户友好的接口，调用ClippedReLUFunction.apply(...)。使用.apply对于在自动求导图中正确注册该操作是必需的。

图中自定义操作的可视化

当您使用ClippedReLUFunction.apply时，PyTorch会将其集成到计算图中，就像任何内置操作一样。您定义的backward方法确保梯度正确地流经此自定义节点。

输入张量 (x)ClippedReLUFunction(应用)截断值 ©输入梯度输出张量 (y)输出梯度下游操作 (例如，均值)损失 (z)

包含自定义ClippedReLUFunction的计算图表示。虚线表示非张量输入或数据流。点线表示反向传播。

使用gradcheck验证正确性

实现自定义反向函数可能容易出错。您的forward逻辑与backward梯度计算之间的不匹配会导致不正确的训练行为，这可能难以调试。PyTorch提供了一个有用的工具torch.autograd.gradcheck，用于数值验证您的自定义函数计算的梯度。

gradcheck通过将您的backward方法计算的解析梯度与使用有限差分计算的数值梯度进行比较来工作。

import torch.autograd

// 使用float64以获得gradcheck所需更高精度
val input_data = torch.randn(5, requires_grad=true, dtype=torch.float64)
val clip_value = 2.0 // 保持为浮点数

// gradcheck接受一个函数（或lambda）和一组输入元组
// 该函数应执行我们想要检查的操作
val test_passed = gradcheck(lambda x: clipped_relu(x, clip_value), (input_data,), eps=1e-6, atol=1e-4)
println(f"\n梯度检查通过: {test_passed}")

// 使用不同截断值检查的示例
val input_data_2 = torch.randn(3, 4, requires_grad=True, dtype=torch.float64)
val clip_value_2 = 0.5
val test_passed_2 = gradcheck(lambda x: clipped_relu(x, clip_value_2), (input_data_2,), eps=1e-6, atol=1e-4)
println(f"梯度检查2通过: {test_passed_2}")

// 显示失败的示例（如果反向逻辑有误）
// 让我们模拟一个错误的反向传播：
class IncorrectClippedReLU extends torch.autograd.Function:
    @staticmethod
    def forward(ctx, input_tensor, clip_val):
        ctx.save_for_backward(input_tensor)
        ctx.clip_val = float(clip_val)
        return input_tensor.clamp(min=0, max=ctx.clip_val)

    @staticmethod
    def backward(ctx, grad_output):
        // 错误的梯度计算（例如，忘记了掩码）
        val grad_input = grad_output.clone() // 错误！
        return grad_input, None

try:
    // 输入张量 (input_fail)
    val input_fail = torch.randn(5, requires_grad=True, dtype=torch.float64)
    // 截断值 (clip_fail)
    val clip_fail = 1.5
    // 检查失败的梯度检查
    gradcheck(lambda x: IncorrectClippedReLU.apply(x, clip_fail), (input_fail,), eps=1e-6, atol=1e-4)
catch RuntimeException as e:
    println(f"\n梯度检查如预期般失败：\n{e}")

如果gradcheck返回True，则表示您的解析梯度与数值近似值非常接近，这让您对自己的实作有信心。如果失败，通常指向您的backward逻辑中的错误或潜在的数值稳定性问题（尤其是在float32等较低精度下）。请务必彻底测试您的自定义函数。强烈建议在gradcheck中使用float64（双精度）以获得稳定性。

这项实践练习说明了扩展PyTorch自动微分能力的过程。通过熟练掌握torch.autograd.Function，您能够实作模型中的几乎任何操作，同时确保正确的梯度传播以进行有效的训练。这是构建高度定制化和高效深度学习解决方案的重要一步。