with torch.no_grad()使用小提示

with torch.no_grad() 是 PyTorch 中的一个上下文管理器，用于禁用梯度计算。当你在模型推理（即生成或预测）时使用它，可以带来以下几个方面的影响和差异：

主要影响

1. 性能提升：

   • 减少内存占用：禁用梯度计算后，PyTorch 不会保存中间变量以供反向传播使用，从而减少了内存的使用。
   • 加速推理速度：由于不需要计算和存储梯度，推理过程中的计算量减少，因此可以加快推理速度。

2. 防止意外更新：

   • 保护模型参数：在推理过程中，通常不希望模型参数被更新。torch.no_grad() 确保了这一点，避免了任何可能的梯度更新操作。

3. 资源优化：

   • 降低显存消耗：对于大型模型，推理时禁用梯度计算可以显著降低显存消耗，使得更多资源可用于其他任务或更大的批量处理。这里需要强调一点，model.generate本身的实现方法中，一般会有装饰器@torch.no_grad()，来表明在调用此方法时将不会计算梯度。如果没有，可以参考如下示例手动添加。

示例代码对比

前提假设： model.generate没有装饰器@torch.no_grad()

使用 torch.no_grad()

with torch.no_grad():
    outputs = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=4096,
        early_stopping=True,
    )

• 优点：
  • 更快的推理速度。
  • 更低的内存占用。
  • 避免不必要的梯度计算和潜在的参数更新。

不使用 torch.no_grad()

outputs = model.generate(
    input_ids,
    attention_mask=attention_mask,
    max_new_tokens=4096,
    early_stopping=True,
)

• 缺点：
  • 推理速度较慢，因为仍然在计算和存储梯度。
  • 内存占用较高，特别是对于大型模型。
  • 可能导致不必要的参数更新，尽管在推理时这通常是不需要的。

实际应用中的选择

• 推理阶段：强烈建议在推理阶段使用 torch.no_grad()，以确保最佳性能和资源利用。
• 训练阶段：在训练阶段，需要计算梯度以更新模型参数，因此不能使用 torch.no_grad()。但在评估（validation）或测试（test）数据集上进行推理时，可以使用 torch.no_grad() 来提高效率。