PyTorch 2.0 反向传播实战：从计算图到梯度下降的 5 步代码实现

Magic Road

307人浏览 · 2026-07-03 13:58:31

Magic Road · 2026-07-03 13:58:31 发布

PyTorch 2.0 反向传播实战：从计算图到梯度下降的 5 步代码实现

深度学习框架的核心竞争力之一，是能够高效地计算梯度并更新模型参数。PyTorch 2.0 通过动态计算图和自动微分机制，将这一过程变得异常简洁。本文将带您从零实现一个完整的反向传播流程，并通过与 PyTorch 原生 Autograd 的对比验证我们的实现。

1. 计算图构建基础

计算图是理解反向传播的视觉化工具。在 PyTorch 中，每个张量操作都会在幕后构建计算图。让我们从一个简单的例子开始：

import torch

# 手动构建计算图
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1
print(f"计算图叶子节点: x={x}, 运算结果: y={y}")

这个简单的例子展示了 PyTorch 如何自动追踪操作历史。当我们在张量上设置 requires_grad=True 时，PyTorch 会记录所有相关操作，构建一个动态计算图。

关键概念 ：

叶子节点 ：计算图的起点（如输入张量）
运算节点 ：对张量执行的操作（如加法、乘法）
梯度传播 ：从输出向输入反向传播梯度

2. 实现基础运算的反向传播

让我们从最基本的加法和乘法运算开始，手动实现它们的反向传播逻辑：

class AddFunction:
    @staticmethod
    def forward(x, y):
        return x + y
    
    @staticmethod
    def backward(grad_output):
        return grad_output, grad_output

class MulFunction:
    @staticmethod
    def forward(x, y):
        return x * y
    
    @staticmethod
    def backward(grad_output, x, y):
        return grad_output * y, grad_output * x

# 测试自定义运算
a = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(4.0, requires_grad=True)

# 前向传播
add_result = AddFunction.forward(a, b)
mul_result = MulFunction.forward(a, b)

print(f"加法结果: {add_result}, 乘法结果: {mul_result}")

这些基础运算构成了更复杂网络的基本构建块。反向传播时，加法运算均匀分配梯度，而乘法运算需要交换输入值进行梯度计算。

3. 构建完整的计算图模块

现在我们将这些基础运算组合成一个可用的计算图模块：

class ComputationGraph:
    def __init__(self):
        self.operations = []
        self.gradients = {}

    def add_operation(self, op_func, inputs):
        output = op_func.forward(*inputs)
        self.operations.append((op_func, inputs, output))
        return output

    def backward(self, initial_grad):
        # 初始化梯度
        self.gradients = {self.operations[-1][2]: initial_grad}
        
        # 反向遍历操作
        for op_func, inputs, output in reversed(self.operations):
            grad_output = self.gradients[output]
            
            if op_func == AddFunction:
                grad_inputs = AddFunction.backward(grad_output)
            elif op_func == MulFunction:
                grad_inputs = MulFunction.backward(grad_output, *inputs)
            
            # 累加梯度（考虑多个输出指向同一输入的情况）
            for input_tensor, grad in zip(inputs, grad_inputs):
                if input_tensor in self.gradients:
                    self.gradients[input_tensor] += grad
                else:
                    self.gradients[input_tensor] = grad

# 使用示例
graph = ComputationGraph()
x = torch.tensor(2.0)
y = torch.tensor(3.0)

a = graph.add_operation(MulFunction, (x, y))
b = graph.add_operation(AddFunction, (a, y))
final_output = graph.add_operation(MulFunction, (b, x))

graph.backward(torch.tensor(1.0))  # 从输出开始反向传播
print(f"x的梯度: {graph.gradients[x]}, y的梯度: {graph.gradients[y]}")

这个实现展示了 PyTorch Autograd 的核心思想——记录操作并在反向传播时应用链式法则。

4. 实现激活函数层

神经网络中的非线性激活函数是模型能够学习复杂模式的关键。让我们实现两个常用的激活函数：

class ReLU:
    @staticmethod
    def forward(x):
        return torch.maximum(torch.tensor(0.0), x)
    
    @staticmethod
    def backward(grad_output, x):
        return grad_output * (x > 0).float()

class Sigmoid:
    @staticmethod
    def forward(x):
        return 1 / (1 + torch.exp(-x))
    
    @staticmethod
    def backward(grad_output, x):
        sigmoid_x = Sigmoid.forward(x)
        return grad_output * sigmoid_x * (1 - sigmoid_x)

# 测试激活函数
test_input = torch.tensor(1.0, requires_grad=True)
relu_out = ReLU.forward(test_input)
sigmoid_out = Sigmoid.forward(test_input)

print(f"ReLU(1.0) = {relu_out}, Sigmoid(1.0) = {sigmoid_out}")

激活函数的反向传播需要特别注意，因为它们通常是非线性变换。ReLU 的梯度在输入大于零时为1，否则为0；而Sigmoid的梯度则可以用其输出值简洁地表示。

5. 完整神经网络实现与验证

现在我们将所有组件组合成一个简单的全连接神经网络，并与 PyTorch 原生实现进行对比验证：

class SimpleNN:
    def __init__(self, input_size, hidden_size, output_size):
        # 初始化权重
        self.W1 = torch.randn(input_size, hidden_size, requires_grad=True)
        self.b1 = torch.randn(hidden_size, requires_grad=True)
        self.W2 = torch.randn(hidden_size, output_size, requires_grad=True)
        self.b2 = torch.randn(output_size, requires_grad=True)
        
    def forward(self, x):
        # 第一层
        z1 = torch.matmul(x, self.W1) + self.b1
        a1 = ReLU.forward(z1)
        
        # 输出层
        z2 = torch.matmul(a1, self.W2) + self.b2
        return z2
    
    def manual_backward(self, x, y_true, loss_fn):
        # 前向传播
        y_pred = self.forward(x)
        loss = loss_fn(y_pred, y_true)
        
        # 反向传播
        grad_loss = torch.tensor(1.0)  # dL/dL = 1
        
        # 输出层梯度
        grad_z2 = grad_loss * (y_pred - y_true)  # 假设使用MSE损失
        grad_W2 = torch.outer(a1, grad_z2)
        grad_b2 = grad_z2
        
        # 隐藏层梯度
        grad_a1 = torch.matmul(grad_z2, self.W2.T)
        grad_z1 = grad_a1 * (z1 > 0).float()  # ReLU导数
        grad_W1 = torch.outer(x, grad_z1)
        grad_b1 = grad_z1
        
        return {
            'W1': grad_W1, 'b1': grad_b1,
            'W2': grad_W2, 'b2': grad_b2
        }

# 对比验证
def verify_gradients():
    # 创建网络和测试数据
    model = SimpleNN(3, 4, 2)
    x = torch.randn(3, requires_grad=True)
    y_true = torch.randn(2)
    
    # 手动计算梯度
    manual_grads = model.manual_backward(x, y_true, lambda y_pred, y: 0.5 * torch.sum((y_pred - y)**2))
    
    # PyTorch自动计算梯度
    y_pred = model.forward(x)
    loss = 0.5 * torch.sum((y_pred - y_true)**2)
    loss.backward()
    
    # 比较结果
    print("手动计算梯度 vs PyTorch自动梯度:")
    for param in ['W1', 'b1', 'W2', 'b2']:
        tensor = getattr(model, param)
        diff = torch.norm(manual_grads[param] - tensor.grad)
        print(f"{param}梯度差异: {diff.item():.6f}")

verify_gradients()

这个完整的实现展示了从输入到输出的完整数据流，以及梯度如何从损失函数传播回网络的第一层。通过与 PyTorch 原生实现的对比，我们可以验证手动实现的正确性。

梯度下降优化实践

理解了反向传播后，我们可以实现一个简单的梯度下降优化器：

class GradientDescentOptimizer:
    def __init__(self, parameters, lr=0.01):
        self.parameters = list(parameters)
        self.lr = lr
    
    def step(self):
        with torch.no_grad():  # 禁用梯度追踪
            for param in self.parameters:
                if param.grad is not None:
                    param -= self.lr * param.grad
    
    def zero_grad(self):
        for param in self.parameters:
            if param.grad is not None:
                param.grad.zero_()

# 训练示例
def train_simple_model():
    # 准备数据
    X = torch.randn(100, 3)
    y = torch.randn(100, 2)
    
    # 初始化模型和优化器
    model = SimpleNN(3, 4, 2)
    optimizer = GradientDescentOptimizer([model.W1, model.b1, model.W2, model.b2], lr=0.01)
    
    # 训练循环
    for epoch in range(100):
        total_loss = 0
        optimizer.zero_grad()
        
        # 前向传播
        y_pred = model.forward(X)
        loss = 0.5 * torch.sum((y_pred - y)**2)
        
        # 反向传播
        loss.backward()
        
        # 参数更新
        optimizer.step()
        
        total_loss += loss.item()
        if epoch % 10 == 0:
            print(f"Epoch {epoch}, Loss: {total_loss:.4f}")

train_simple_model()

这个训练循环展示了深度学习中的关键步骤：前向传播、损失计算、反向传播和参数更新。PyTorch 的自动微分系统让我们能够专注于模型设计，而无需手动计算复杂的梯度。

性能优化技巧

PyTorch 2.0 引入了多项性能优化，我们可以借鉴这些思想来改进我们的实现：

算子融合 ：将多个小操作合并为一个大操作
内存优化 ：重用中间结果的内存
并行计算 ：利用GPU的并行能力

# 优化后的矩阵乘法实现
def optimized_matmul(A, B):
    # 实际项目中会使用更高效的实现
    return torch.matmul(A, B)

# 使用@torch.compile装饰器加速(PyTorch 2.0特性)
@torch.compile
def fast_forward(x, W1, b1, W2, b2):
    z1 = torch.matmul(x, W1) + b1
    a1 = torch.relu(z1)
    z2 = torch.matmul(a1, W2) + b2
    return z2

理解这些底层原理对于调试模型和实现自定义操作至关重要。当遇到性能瓶颈或需要特殊操作时，这些知识将非常有用。