PyTorch `torch.no_grad` vs `torch.inference_mode`

在 PyTorch 中，torch.no_grad 和 torch.inference_mode 都用于在推理（inference）过程中禁用梯度计算，以提高性能和减少内存消耗。然而，它们之间有一些关键的区别和使用场景。

torch.no_grad

torch.no_grad 是一个上下文管理器，用于临时禁用梯度计算。它通常用于推理阶段，以确保在前向传播过程中不计算梯度，从而节省内存和计算资源。

示例用法

import torch

model = ...  # 你的模型
input_tensor = ...  # 输入张量

with torch.no_grad():
    output = model(input_tensor)

特点

1. 临时禁用梯度计算：在 with 语句块内禁用梯度计算，块外恢复正常。
2. 适用于推理阶段：主要用于推理阶段，以提高性能和减少内存消耗。
3. 兼容性好：与大多数现有代码兼容，不需要对代码进行大的修改。

torch.inference_mode

torch.inference_mode 是 PyTorch 1.9.0 引入的一个新的上下文管理器，专门用于推理阶段。与 torch.no_grad 类似，它也禁用梯度计算，但它还做了更多优化，以进一步提高性能和减少内存消耗。

示例用法

import torch

model = ...  # 你的模型
input_tensor = ...  # 输入张量

with torch.inference_mode():
    output = model(input_tensor)

特点

1. 更高的性能优化：除了禁用梯度计算外，还进行了其他优化，以进一步提高推理性能。
2. 适用于推理阶段：专门为推理阶段设计，提供更高效的推理性能。
3. 更严格的限制：在 torch.inference_mode 下，某些操作可能会被限制，以确保性能优化。例如，某些需要梯度计算的操作可能会被禁止。

选择使用哪一个

• torch.no_grad：如果你需要在推理阶段禁用梯度计算，并且希望代码兼容性更好，使用 torch.no_grad 是一个不错的选择。
• torch.inference_mode：如果你希望在推理阶段获得更高的性能，并且可以接受某些操作的限制，使用 torch.inference_mode 是更好的选择。

性能对比

torch.inference_mode 通常比 torch.no_grad 提供更高的性能优化，因为它不仅禁用梯度计算，还进行了其他优化。然而，这些优化可能会带来一些限制，因此在选择使用哪一个时需要根据具体需求进行权衡。

总结

• torch.no_grad：适用于需要临时禁用梯度计算的场景，兼容性好，适用于大多数推理任务。
• torch.inference_mode：适用于需要更高推理性能的场景，提供更高效的推理性能，但可能会有一些操作限制。