Published on

深入理解 PyTorch 中的 `with torch.no_grad()`

在 Python 中，with 语句是一种上下文管理器（Context Manager），用于简化资源的分配和释放过程。它确保在代码块执行前初始化资源，执行后自动清理资源，避免手动管理资源可能导致的错误（如忘记关闭文件、网络连接或释放锁）。

基本语法与工作原理

with 表达式 [as 变量]:
    代码块
  1. 表达式：返回一个上下文管理器对象（需实现 __enter__() 和 __exit__() 方法）。
  2. __enter__()：进入上下文时调用，返回值可赋给 as 后的变量（可选）。
  3. 代码块：执行具体操作。
  4. __exit__()：退出上下文时自动调用，负责清理资源（如关闭文件）。

在 PyTorch 中的应用场景

在你的代码中，with torch.no_grad(): 是 PyTorch 特有的上下文管理器，用于临时禁用梯度计算，主要有以下用途：

1. 推理阶段加速计算

在模型推理时，无需计算梯度，可通过 torch.no_grad() 减少内存消耗并提升计算速度。

with torch.no_grad():
    outputs = model(inputs)  # 前向传播但不跟踪梯度

2. 手动更新参数（如SGD）

你的代码片段正是这种场景：

with torch.no_grad():
    params -= learning_rate * params.grad  # 手动梯度更新
  • 为什么用 with torch.no_grad()：
    PyTorch 默认会跟踪所有张量操作的梯度（若 requires_grad=True）。手动更新参数时，我们不希望这些更新操作被记录为计算图的一部分（否则会导致梯度计算错误或内存泄漏）。torch.no_grad() 临时禁用梯度计算，确保参数更新不影响后续的反向传播。

等价的手动实现方式

若不使用 with 语句，需手动管理梯度计算状态：

# 禁用梯度
torch.set_grad_enabled(False)
try:
    params -= learning_rate * params.grad
finally:
    # 恢复梯度
    torch.set_grad_enabled(True)

显然，with 语句使代码更简洁、安全。

其他常见的 PyTorch 上下文管理器

  • torch.enable_grad()：强制启用梯度计算（即使在 no_grad() 嵌套中）。
  • torch.inference_mode()：比 no_grad() 更轻量级的推理模式，进一步减少内存占用。
  • torch.autocast()：用于混合精度训练，自动处理浮点精度转换。

总结

with torch.no_grad(): 的核心作用是临时关闭梯度计算，常用于：

  1. 推理阶段提升性能；
  2. 手动参数更新（如自定义优化器）；
  3. 减少不必要的内存占用（如计算验证集指标时）。

它是 PyTorch 中高效、安全地执行无梯度操作的标准方式。

THE END