自动微分机制深度剖析
PyTorch 的自动微分系统(autograd)是深度学习框架的核心。理解其计算图构建和反向传播机制,有助于调试梯度问题和实现自定义算子。
计算图构建
python
# 动态计算图示例
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
z = y + 3 * x
# 反向传播
z.backward()
print(x.grad) # tensor([7.]) = 2x + 3 = 2*2 + 3
# 计算图结构
# x → (x²) → y → (+ 3x) → z自定义梯度函数
python
class CustomReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input梯度检查
使用 torch.autograd.gradcheck 验证自定义梯度函数的正确性。它通过数值微分与解析梯度对比,确保实现无误。