深入理解 PyTorch 自动微分机制与自定义 torch.autograd.Function

---

前言

随着深度学习技术的迅速发展，PyTorch 作为一款功能强大且灵活的深度学习框架，受到了广泛的关注和应用。它以其动态计算图、易用性以及强大的社区支持而闻名。在PyTorch中，自动微分（autograd）是其核心特性之一，它使得神经网络训练过程中的梯度计算变得简单高效。对于大多数应用场景而言，开发者无需手动编写反向传播逻辑，因为PyTorch能够自动处理这些细节。

然而，在某些特殊情况下，我们可能需要对特定的操作进行定制化的梯度计算，这时就需要深入了解并利用PyTorch提供的torch.autograd.Function类来实现自定义的前向和后向传播逻辑。通过这种方式，不仅可以实现更复杂的模型结构，还能优化性能或满足特定的研究需求。

本文将从基础出发，首先介绍如何使用PyTorch内置的自动微分机制完成常规的模型训练流程；接着详细解析torch.autograd.Function中的ctx对象及其在前后向传播间的作用；最后，通过一个具体的例子演示如何编写自定义的自动微分函数，并解释其中的关键概念和操作。希望通过这篇文章，读者能够掌握PyTorch自动微分的核心原理，以及如何根据实际需求设计高效的自定义梯度计算逻辑。

一、pytorch使用现有的自动微分机制

编写一个后向传播函数在 PyTorch 中通常是不需要的，因为 PyTorch 自动处理了自动微分（autograd），即通过 loss.backward() 来计算梯度。下面我们将展示如何编写一个简单的自定义后向传播函数，并解释如何在 PyTorch 中利用现有的自动微分机制进行反向传播。

通常情况下，你只需要调用 loss.backward() 即可完成反向传播，一个示列代码如下：

import torch
import torch.nn as nn
import torch.optim as optim# 定义一个简单的模型
class SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()self.fc = nn.Linear(10, 1)def forward(self, x):return self.fc(x)# 初始化模型、损失函数和优化器
model = SimpleModel()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)# 输入和目标
input_tensor = torch.randn(1, 10, requires_grad=True)
target = torch.tensor([[1.]])# 前向传播
output = model(input_tensor)
loss = criterion(output, target)# 清除之前的梯度
optimizer.zero_grad()# 反向传播
loss.backward()# 更新参数
optimizer.step()

二、torch.autograd.Function中的ctx解读

在PyTorch中，torch.autograd.Function 是用来定义自定义自动求导函数的类。你提供的CustomReLU类继承了torch.autograd.Function并实现了自定义的前向传播和反向传播逻辑。这里的ctx（context）对象是用于存储信息以便在前向传播和反向传播之间共享。

1、forward 方法中的 ctx

在forward方法中，ctx被用来保存在前向传播阶段计算的信息，这些信息可能在后续的反向传播过程中需要使用。例如：

@staticmethod
def forward(ctx, input):ctx.save_for_backward(input)  # 保存输入以供反向传播使用return input.clamp(min=0)

• ctx.save_for_backward(input)：这里我们保存了输入张量input。这很重要，因为在反向传播时我们需要知道哪些元素在前向传播中被设为零（即负数），以便正确地将梯度设为零。

2、backward 方法中的 ctx

在backward方法中，ctx被用来访问在前向传播阶段保存的信息。例如：

@staticmethod
def backward(ctx, grad_output):input, = ctx.saved_tensorsgrad_input = grad_output.clone()grad_input[input < 0] = 0return grad_input

• input, = ctx.saved_tensors：这里我们从ctx中获取了在前向传播阶段保存的输入张量。注意，saved_tensors是一个元组，即使只保存了一个张量，也需要用逗号来解包。

• 接下来，我们基于原始输入创建了grad_input，它初始化为grad_output的副本。然后我们将所有在前向传播中对应的输入小于0的位置的梯度设为0，这是因为ReLU激活函数对于所有负值输入都输出0，所以其梯度也应为0。

3、小结

ctx的作用是在前向传播和反向传播之间传递必要的信息。通过ctx.save_for_backward()可以在前向传播中保存任何需要在反向传播中使用的数据，而在反向传播中则可以通过ctx.saved_tensors来访问这些数据。这对于实现自定义的自动求导函数来说是非常重要的，因为它允许我们在不需要显式管理复杂状态的情况下执行复杂的梯度计算。

三、pytorch自定义自动微分函数（torch.autograd.Function）

1、torch.autograd.Function计算前向与后向传播梯度Demo

然而，如果你确实需要自定义反向传播逻辑或理解其工作原理，可以通过定义自定义的自动微分函数来实现。如果你想自定义某些操作的反向传播逻辑，可以使用 torch.autograd.Function 来创建自定义的自动微分函数。以下是一个简单的例子：

import torchclass CustomReLU(torch.autograd.Function):@staticmethoddef forward(ctx, input):"""在前向传播中，我们接收到一个上下文对象和一个输入张量，并返回一个经过 ReLU 激活的输出张量。"""ctx.save_for_backward(input)  # 保存输入以供反向传播使用return input.clamp(min=0)@staticmethoddef backward(ctx, grad_output):"""在反向传播中，我们接收到一个上下文对象和一个输出张量的梯度，并返回输入张量的梯度。"""input, = ctx.saved_tensorsgrad_input = grad_output.clone()grad_input[input < 0] = 0return grad_input# 使用自定义 ReLU 函数
custom_relu = CustomReLU.apply# 示例：将自定义 ReLU 应用于输入
input_tensor = torch.randn(5, requires_grad=True)
output = custom_relu(input_tensor)# 创建一个简单的损失并进行反向传播
loss = output.sum()
loss.backward()print("Input tensor:", input_tensor)
print("Gradient of input tensor:", input_tensor.grad)

2、前向传播梯度解读

如果要计算前向传播梯度只有执行代码output = custom_relu(input_tensor)才能启动（input_tensor是一个维度[ …]），因此当启动了代码，就可以执行我们定义前向传播方法：

@staticmethod
def forward(ctx, input):"""在前向传播中，我们接收到一个上下文对象和一个输入张量，并返回一个经过 ReLU 激活的输出张量。"""ctx.save_for_backward(input)  # 保存输入以供反向传播使用return input.clamp(min=0)

而ctx.save_for_backward也是保存输入input内容，custom_relu只有一层模型，就是来自input_tensor值，所以ctx保存了input_tensor值，也是一个维度值。

3、后向传播梯度解读

如果要计算后向传播梯度只有执行代码loss.backward()才能启动，因此当启动了代码，就可以执行我们定义后向传播方法：

@staticmethod
def backward(ctx, grad_output):"""在反向传播中，我们接收到一个上下文对象和一个输出张量的梯度，并返回输入张量的梯度。"""input, = ctx.saved_tensorsgrad_input = grad_output.clone()grad_input[input < 0] = 0return grad_input

而ctx.saved_tensors是取前向保存的内容。

4、运行结果