PyTorch源码解析：torch.no_grad()是怎么巧妙实现的

PyTorch作为一种流行的机器学习框架，提供了许多便利的功能，以协助开发者编写高效和优雅的代码。其中之一便是torch.no_grad()函数，它通过禁用梯度计算来提高推理效率。本文将深入探讨torch.no_grad()的源码实现，并重点介绍了其与with语句和Python装饰器的妙用。

一、torch.no_grad()的基本介绍在深度学习中，梯度计算是训练过程中的关键步骤，但在推理过程中并不需要计算梯度。为了提高推理效率，PyTorch提供了torch.no_grad()函数，它可以在推理过程中临时禁用梯度计算。这样一来，PyTorch将跳过梯度计算的相关操作，从而减少了内存消耗和计算开销。

二、torch.no_grad()的源码实现

下面是torch.no_grad()函数的源码实现：

def no_grad():
    """    Context-manager that disabled gradient calculation.    Disabling gradient calculation is useful for inference, when you are sure that you will not call :meth:`Tensor.backward()`.    It will reduce memory consumption for computations that would otherwise have `requires_grad=True`.    In this mode, the result of every computation will have `requires_grad=False`, even when the inputs have `requires_grad=True`.    This context manager is thread local; it will not affect computation in other threads.    Also functions as a decorator. (Make sure to instantiate with parentheses.)    Example::        >>> x = torch.tensor([1], requires_grad=True)        >>> with torch.no_grad():        ...   y = x * 2        >>> y.requires_grad        False    """
    return _DecoratorContextManager(_no_grad)

可以看到，torch.no_grad()实际上返回了一个名为_DecoratorContextManager的装饰器上下文管理器。这个上下文管理器在内部调用了_no_grad函数，并返回一个包装后的计算结果。

_no_grad函数的实现如下：

def _no_grad():
    is_grad_enabled = torch.is_grad_enabled()
    try:
        torch.set_grad_enabled(False)
        yield
    finally:
        torch.set_grad_enabled(is_grad_enabled)

_no_grad函数第一使用torch.is_grad_enabled()函数获取当前的梯度计算开关状态，并保存在is_grad_enabled变量中。然后，它调用torch.set_grad_enabled(False)将梯度计算开关设置为False，从而禁用梯度计算。接下来，yield语句用于将控制权交给包装后的计算代码块。最后，在finally块中，它调用torch.set_grad_enabled(is_grad_enabled)将梯度计算开关恢复到之前的状态。

三、torch.no_grad()与with语句的妙用torch.no_grad()函数可以与with语句结合使用，以提供一个临时禁用梯度计算的代码块。在with torch.no_grad():的上下文中，所有的计算操作将自动禁用梯度计算，从而提高推理阶段的性能。这种结构和语法的设计超级巧妙，使得代码编写更加简洁和可读。

例如，下面是一个使用torch.no_grad()和with语句的示例：

import torch

x = torch.tensor([1], requires_grad=True)
w = torch.tensor([2], requires_grad=True)

with torch.no_grad():
    y = x * w

print(y.requires_grad)  # 输出: False

在上述示例中，with torch.no_grad():语句块内的计算操作y = x * w会自动禁用梯度计算，所以y的requires_grad属性被设置为False。这意味着在这个上下文中，y不会参与梯度传播，从而减少了内存消耗和计算开销。

四、torch.no_grad()与Python装饰器的妙用除了可以用作with语句的上下文管理器外，torch.no_grad()还可以作为Python装饰器使用。装饰器是一种Python语法，用于在不修改原始函数代码的情况下，给函数添加额外的功能。

下面是一个使用torch.no_grad()作为装饰器的示例：

import torch

@torch.no_grad()
def inference(x):
    w = torch.tensor([2], requires_grad=True)
    y = x * w
    return y

x = torch.tensor([1])
y = inference(x)

print(y.requires_grad)  # 输出: False

在上述示例中，@torch.no_grad()装饰器被应用于inference函数。这意味着当调用inference函数时，函数内部的计算操作将自动禁用梯度计算。因此，返回的y张量的requires_grad属性被设置为False。

通过使用torch.no_grad()装饰器，我们可以轻松地将梯度计算的禁用功能应用于函数，而无需在函数内部显式添加with torch.no_grad():语句块。

总结：本文深入介绍了PyTorch中torch.no_grad()函数的源码实现，并重点突出了它与with语句和Python装饰器的妙用。torch.no_grad()的设计使得禁用梯度计算变得简单而优雅，提高了推理阶段的性能和代码的可读性。通过充分理解和运用torch.no_grad()的特性，我们可以更好地利用PyTorch框架提供的功能，编写高效且可维护的深度学习代码。

摘自：科学随想录