PyTorch/TensorFlow 自定义 `autograd`：实现复杂梯度的自动求导

好的，各位观众，欢迎来到“PyTorch/TensorFlow 自定义 autograd：实现复杂梯度的自动求导”讲座现场！今天咱们就来聊聊深度学习框架里一个非常酷炫的功能——自定义 autograd。这玩意儿就像是给框架装了个 turbo 引擎，让它能处理更复杂、更个性化的梯度计算。

第一部分：什么是 Autograd？为啥我们需要自定义它？

首先，咱们得搞清楚 autograd 是个啥。简单来说，autograd 就是“自动求导”的意思。它负责追踪你的张量（Tensor）运算，然后自动帮你计算梯度。这就像你辛辛苦苦写了一个复杂的数学公式，然后有个小精灵自动帮你算出每个变量的导数，简直不要太爽！

PyTorch 和 TensorFlow 都内置了强大的 autograd 引擎，能够处理大部分常见的操作。但是！人生总有意外，总有一些场景是内置的 autograd 搞不定的。比如：

非标准操作： 你自己发明了一个新的激活函数，或者一个新的损失函数，框架里没有现成的梯度计算公式。
性能优化： 某些操作的梯度计算非常耗时，你想用更高效的算法来加速它。
特定需求： 你需要在梯度计算过程中加入一些特殊的逻辑，比如梯度裁剪、梯度扰动等等。
遗留代码对接: 你的代码里面有一些使用numpy或者其他科学计算库的函数，你需要将它们融入到PyTorch或者TensorFlow的计算图中

这时候，自定义 autograd 就显得非常重要了。它可以让你突破框架的限制，实现更灵活、更强大的梯度计算。

第二部分：PyTorch 自定义 Autograd 的正确姿势

在 PyTorch 中，自定义 autograd 主要涉及以下几个步骤：

定义一个继承自 torch.autograd.Function 的类。 这个类就是你自定义操作的蓝图。
实现 forward 方法。 forward 方法定义了前向计算的逻辑，也就是根据输入计算输出的过程。
实现 backward 方法。 backward 方法定义了反向传播的逻辑，也就是根据输出的梯度计算输入的梯度的过程。

下面咱们来举个栗子，假设我们要自定义一个叫做 MyReLU 的 ReLU 激活函数，但是这次我们想对梯度进行一些修改，当输入小于0时，梯度变为0.1。

import torch

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        """
        前向传播：计算 ReLU 输出。
        ctx 是一个上下文对象，可以用来保存需要在 backward 中使用的变量。
        """
        ctx.save_for_backward(input)  # 保存输入，以便在 backward 中使用
        return input.clamp(min=0)  # ReLU 的 forward 就是把小于 0 的值变成 0

    @staticmethod
    def backward(ctx, grad_output):
        """
        反向传播：计算输入梯度。
        grad_output 是输出的梯度。
        """
        input, = ctx.saved_tensors  # 取出 forward 中保存的输入
        grad_input = grad_output.clone()  # 创建一个输出梯度的副本
        grad_input[input < 0] = 0.1  # 根据 ReLU 的梯度公式计算输入梯度
        return grad_input

# 使用自定义的 ReLU
my_relu = MyReLU.apply  # 注意这里要用 apply 方法
x = torch.randn(10, requires_grad=True)
y = my_relu(x)
z = y.sum()
z.backward()

print("输入 x 的梯度：", x.grad)

咱们来解释一下这段代码：

MyReLU 类继承了 torch.autograd.Function。
forward 方法接收一个输入 input，并使用 input.clamp(min=0) 计算 ReLU 输出。ctx.save_for_backward(input) 这行代码非常重要，它把 input 保存起来，以便在 backward 方法中使用。
backward 方法接收一个输出梯度 grad_output，并根据 ReLU 的梯度公式计算输入梯度。input, = ctx.saved_tensors 从 ctx 中取出之前保存的 input。grad_input[input < 0] = 0.1 这行代码就是 ReLU 的梯度公式：当输入小于 0 时，梯度为0.1。
最后，我们使用 MyReLU.apply 来调用自定义的 ReLU。

自定义 Autograd 的进阶技巧

保存多个变量： ctx.save_for_backward 可以保存多个变量，用逗号分隔即可。
处理非张量输入： forward 方法可以接收非张量输入，但 backward 方法只能返回张量。
使用 torch.no_grad： 在 forward 和 backward 方法中，可以使用 torch.no_grad 来禁用梯度计算，提高性能。

第三部分：TensorFlow 自定义 Autograd 的正确姿势

在 TensorFlow 中，自定义 autograd 稍微复杂一些，主要涉及以下几个步骤：

使用 tf.custom_gradient 装饰器。 这个装饰器可以让你自定义一个函数的梯度计算逻辑。
定义一个函数，接收输入和梯度，并返回计算后的梯度。 这个函数就是你自定义的梯度计算逻辑。

下面咱们也来举个栗子，还是自定义一个叫做 MyReLU 的 ReLU 激活函数，并且对梯度进行修改。

import tensorflow as tf

@tf.custom_gradient
def my_relu(x):
    """
    自定义 ReLU 激活函数。
    """
    def grad(dy):
        """
        自定义梯度计算函数。
        dy 是输出的梯度。
        """
        return dy * tf.cast(x > 0, tf.float32) + 0.1 * tf.cast(x <= 0, tf.float32)  # 根据 ReLU 的梯度公式计算输入梯度

    return tf.nn.relu(x), grad  # 返回前向计算结果和自定义的梯度计算函数

# 使用自定义的 ReLU
x = tf.Variable(tf.random.normal((10,)))
with tf.GradientTape() as tape:
    y = my_relu(x)
z = tf.reduce_sum(y)
grad = tape.gradient(z, x)

print("输入 x 的梯度：", grad)

咱们来解释一下这段代码：

@tf.custom_gradient 装饰器告诉 TensorFlow，my_relu 函数需要自定义梯度计算。
grad 函数接收一个输出梯度 dy，并根据 ReLU 的梯度公式计算输入梯度。dy * tf.cast(x > 0, tf.float32) 这部分计算了 ReLU 在 x > 0 时的梯度，0.1 * tf.cast(x <= 0, tf.float32) 计算了 ReLU 在 x <= 0 时的梯度。
my_relu 函数返回前向计算结果 tf.nn.relu(x) 和自定义的梯度计算函数 grad。
最后，我们使用 tf.GradientTape 来计算梯度。

自定义 Autograd 的进阶技巧

处理多个输入： grad 函数可以接收多个输入，对应于 my_relu 函数的多个输入。
返回多个梯度： grad 函数可以返回多个梯度，对应于 my_relu 函数的多个输入。
使用 tf.stop_gradient： 在 grad 函数中，可以使用 tf.stop_gradient 来阻止某些变量的梯度计算。

第四部分：自定义 Autograd 的应用场景

自定义 autograd 的应用场景非常广泛，这里列举一些常见的例子：

自定义激活函数： 可以自定义各种各样的激活函数，比如 Swish、Mish 等等。
自定义损失函数： 可以自定义各种各样的损失函数，比如 Focal Loss、Dice Loss 等等。
梯度裁剪： 可以对梯度进行裁剪，防止梯度爆炸。
梯度扰动： 可以对梯度进行扰动，提高模型的泛化能力。
对抗训练： 可以使用自定义 autograd 来实现对抗训练，提高模型的鲁棒性。
优化算法定制： 可以将一些传统的优化算法融入到深度学习框架中，比如L-BFGS

表格：PyTorch vs. TensorFlow 自定义 Autograd

特性	PyTorch	TensorFlow
主要方式	继承 `torch.autograd.Function`	使用 `tf.custom_gradient` 装饰器
前向传播	`forward` 方法	函数本身
反向传播	`backward` 方法	`grad` 函数
上下文对象	`ctx`	无，通过闭包传递变量
保存变量	`ctx.save_for_backward`	通过闭包保存
禁用梯度计算	`torch.no_grad`	`tf.stop_gradient`
灵活性	较高，更接近底层实现	较高，但需要理解函数式编程和闭包
易用性	相对容易理解，更符合面向对象编程习惯	相对复杂，需要理解装饰器和函数式编程

第五部分：自定义 Autograd 的注意事项

性能问题： 自定义 autograd 可能会引入额外的性能开销，需要仔细评估。
梯度检查： 自定义 autograd 很容易出错，需要使用梯度检查来验证梯度的正确性。
数值稳定性： 自定义 autograd 可能会引入数值不稳定性，需要仔细处理。

梯度检查

梯度检查是一种验证自定义 autograd 实现是否正确的方法。它的基本思想是：使用有限差分法来近似计算梯度，然后将近似梯度与自定义 autograd 计算的梯度进行比较。如果两者之间的误差很小，说明自定义 autograd 实现是正确的。

PyTorch 梯度检查示例：

import torch
from torch.autograd import gradcheck

# 定义一个简单的自定义操作
class MyOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * x

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        return 2 * x * grad_output

# 创建一个输入张量
x = torch.randn(10, requires_grad=True)

# 使用 gradcheck 进行梯度检查
test = gradcheck(MyOp.apply, (x,), eps=1e-6, atol=1e-4)
print(test)

TensorFlow 梯度检查示例：

import tensorflow as tf
import numpy as np

@tf.custom_gradient
def my_func(x):
    def grad(dy):
        return dy * tf.cast(x > 0, tf.float32)
    return tf.nn.relu(x), grad

# 创建一个输入张量
x = tf.Variable(tf.random.normal((10,)))

# 使用数值方法计算梯度
def compute_numerical_gradients(func, input_tensor, delta=1e-6):
    grads = []
    for i in range(input_tensor.shape[0]):
        original_value = input_tensor[i].numpy()
        # 计算正向扰动
        input_tensor[i].assign(original_value + delta)
        with tf.GradientTape() as tape:
            y_plus = func(input_tensor)
        grad_plus = tape.gradient(y_plus, input_tensor)
        # 计算反向扰动
        input_tensor[i].assign(original_value - delta)
        with tf.GradientTape() as tape:
            y_minus = func(input_tensor)
        grad_minus = tape.gradient(y_minus, input_tensor)
        # 计算数值梯度
        numerical_grad = (grad_plus[i].numpy() - grad_minus[i].numpy()) / (2 * delta)
        grads.append(numerical_grad)
        # 恢复原始值
        input_tensor[i].assign(original_value)
    return np.array(grads)

# 使用 GradientTape 计算梯度
with tf.GradientTape() as tape:
    y = my_func(x)
analytical_grads = tape.gradient(y, x).numpy()

# 计算数值梯度
numerical_grads = compute_numerical_gradients(my_func, x)

# 比较两者之间的误差
error = np.mean(np.abs(analytical_grads - numerical_grads))
print("Error:", error)

第六部分：总结

好了，各位观众，今天的“PyTorch/TensorFlow 自定义 autograd：实现复杂梯度的自动求导”讲座就到这里了。希望通过今天的讲解，大家能够掌握自定义 autograd 的基本方法和技巧，并能够灵活地应用到自己的项目中。

记住，自定义 autograd 是一把双刃剑，它可以让你突破框架的限制，实现更强大的功能，但也可能会引入额外的性能开销和错误。因此，在使用自定义 autograd 时，一定要谨慎评估，并进行充分的测试和验证。

最后，祝大家编程愉快，bug 远离！下次再见！

发表回复 取消回复

发表回复取消回复