【PyTorch】torch.no_grad() 函数：禁用自动求导（Autograd）的功能

torch.no_grad() 函数

torch.no_grad() 是 PyTorch 提供的一个上下文管理器（context manager），用于在其上下文块中禁用自动求导（Autograd）的功能，从而避免梯度计算和反向传播。这种方式可以显著减少内存消耗，加快推理过程，因此通常用于模型的评估或推理阶段。

1. 功能和特点

• 禁用梯度计算：在 torch.no_grad() 块内，所有涉及 requires_grad=True 张量的操作都不会生成计算图，也不会存储中间梯度。
• 减少内存占用：由于不需要存储中间结果用于反向传播，内存使用显著降低。
• 加速推理：避免了多余的计算，尤其在大模型的推理过程中，能提高性能。

2. 语法

with torch.no_grad():
    # 你的代码块

或者作为装饰器使用：

@torch.no_grad()
def inference(model, data):
    # 推理代码

3. 使用场景

3.1 模型评估和推理

在模型评估或推理阶段，我们只关心模型的前向传播结果，不需要计算梯度，因此 torch.no_grad() 是标准用法。

import torch
import torch.nn as nn

# 定义一个简单模型
model = nn.Linear(10, 1)
model.eval()  # 切换到评估模式

# 输入数据
inputs = torch.randn(5, 10)

# 使用 torch.no_grad() 进行推理
with torch.no_grad():
    outputs = model(inputs)
    print(outputs)

3.2 避免无意的梯度计算

在某些场景下，我们可能会对 requires_grad=True 的张量进行额外的运算，但不希望这些操作影响梯度计算。这种情况下可以使用 torch.no_grad()。

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

with torch.no_grad():
    y = x * 2  # 不会记录梯度
print(y.requires_grad)  # False

3.3 减少内存占用

在不需要梯度的情况下，通过禁用梯度计算，可以显著减少内存占用。这在大型模型推理中尤其重要。

4. 注意事项

1. torch.no_grad() 的作用域

   • torch.no_grad() 的作用范围只限于其上下文块内，退出块后会恢复默认的 Autograd 功能。

   x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
   
   with torch.no_grad():
       y = x * 2
   z = x * 3
   
   print(y.requires_grad)  # False
   print(z.requires_grad)  # True
   

2. 与 model.eval() 的区别

   • torch.no_grad()：禁用梯度计算，不影响模型的运行模式。
   • model.eval()：切换模型到评估模式（如禁用 Dropout、BatchNorm 的训练行为），但不会禁用梯度计算。

   在评估或推理时，通常会同时使用：

   model.eval()
   with torch.no_grad():
       outputs = model(inputs)
   

3. 影响计算结果的层

   • torch.no_grad() 不影响模型的前向传播结果，仅影响梯度计算，因此它不会改变模型的计算逻辑或输出。

5. 实际例子

5.1 对比推理性能

import torch
import torch.nn as nn
import time

# 定义模型和数据
model = nn.Linear(1000, 1000).cuda()
inputs = torch.randn(1000, 1000).cuda()

# 使用梯度计算的推理
start = time.time()
for _ in range(100):
    outputs = model(inputs)
end = time.time()
print("With grad:", end - start)

# 使用 torch.no_grad() 的推理
start = time.time()
with torch.no_grad():
    for _ in range(100):
        outputs = model(inputs)
end = time.time()
print("Without grad:", end - start)

5.2 自定义推理函数

@torch.no_grad()
def inference(model, inputs):
    model.eval()
    return model(inputs)

# 调用推理函数
outputs = inference(model, inputs)

6. 总结

• torch.no_grad() 是禁用自动梯度求导的一种方式，常用于模型推理或评估。
• 它可以减少内存占用、提高运行效率。
• 与 model.eval() 配合使用时，能够实现更高效的推理。
• 它不会改变模型的计算逻辑，仅仅影响梯度计算和资源占用。

适当使用 torch.no_grad() 能显著提高代码性能，特别是在深度学习模型推理阶段。