PyTorch 基础学习（6）- 自动微分

系列文章：
PyTorch 基础学习（1） - 快速入门
PyTorch 基础学习（2）- 张量 Tensors
PyTorch 基础学习（3） - 张量的数学操作
PyTorch 基础学习（4）- 张量的类型
PyTorch 基础学习（5）- 神经网络
PyTorch 基础学习（6）- 函数API

自动微分的作用

在深度学习和机器学习中，自动微分是用于计算梯度的核心技术。梯度在优化过程中起着至关重要的作用，它们用于调整模型参数以最小化损失函数。通过计算模型输出相对于输入参数的梯度，自动微分可以帮助我们在反向传播算法中高效地更新模型参数。

自动微分工具包的应用场景

torch.autograd 是 PyTorch 中用于自动微分的工具包，它主要应用于以下场景：

1. 神经网络训练：在反向传播过程中自动计算损失函数相对于模型参数的梯度。
2. 优化算法：使用梯度信息来优化模型参数，使损失函数达到极小值。
3. 复杂计算图的求导：支持对任意复杂计算图的自动求导，方便进行自定义模型和损失函数的优化。

自动求导的基本使用

使用 torch.autograd.backward 函数

import torch
from torch.autograd import Variable

# 创建一个张量并包装在Variable中
x = Variable(torch.tensor([1.0, 2.0, 3.0]), requires_grad=True)

# 定义一个简单的标量函数
y = x.sum()

# 计算梯度
y.backward()

# 输出梯度
print(x.grad)  # 输出: tensor([1., 1., 1.])

在上面的代码中，我们定义了一个标量函数 y = x.sum()，并使用 y.backward() 计算梯度。由于 y 是一个标量，因此我们不需要为 backward() 函数提供额外的参数。

处理非标量输出

如果输出不是标量（例如一个向量），我们需要在调用 backward() 时提供 grad_variables 参数，该参数用于指定与输出同形状的梯度。

z = x * 2

# 指定每个输出元素的梯度
z.backward(torch.tensor([1.0, 1.0, 1.0]))

print(x.grad)  # 输出: tensor([3., 3., 3.]) (之前的梯度已累积)

在这个例子中，我们计算了 z = x * 2 的梯度，并指定了一个与 z 同形状的梯度 [1.0, 1.0, 1.0]，使得每个元素都被正确求导。

使用 retain_graph 参数

当我们需要对同一个计算图多次进行求导时，可以使用 retain_graph=True 来保留中间计算结果：

y = x.sum()
y.backward(retain_graph=True)
y.backward()  # 可以多次调用

Variable 对象

Variable 是 Tensor 的一个封装，记录了对其进行的操作。对于在 Variable 上进行的每一个操作，autograd 都会记录操作历史，从而支持反向传播。

在 PyTorch 0.4 及以后的版本中，Variable 已经合并到 Tensor 中，默认的 Tensor 就支持自动求导。

钩子（Hooks）

可以在 Variable 上注册钩子函数，每次计算梯度时都会调用这些钩子。钩子不应修改输入，但可以返回新的梯度。

def hook_fn(grad):
    return grad * 2

v = Variable(torch.tensor([0.0, 0.0, 0.0]), requires_grad=True)
h = v.register_hook(hook_fn)

v.backward(torch.tensor([1.0, 1.0, 1.0]))
print(v.grad)  # 输出: tensor([2., 2., 2.])

h.remove()  # 移除钩子

在这个例子中，我们注册了一个钩子函数，该函数将梯度值加倍。

强化学习中的奖励

在涉及随机操作的计算图中，必须在随机节点的输出上调用 reinforce() 方法来提供奖励值：

reward = torch.tensor([0.5, 0.5, 0.5])
v.reinforce(reward)

创建自定义 Function

自定义的 Function 可以通过子类化 torch.autograd.Function 来定义。用户需要实现 forward 和 backward 方法。

class MyReLU(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

在这个自定义 ReLU 函数中，我们在 forward 方法中使用 ctx.save_for_backward() 保存输入张量，并在 backward 方法中通过 ctx.saved_tensors 获取它来计算梯度。

应用实例：多层前馈神经网络

下面是一个完整的应用实例脚本，该脚本展示了如何使用 torch.autograd 进行自动微分，并结合了教程中的多个元素，包括自定义函数、梯度计算、钩子、和使用 retain_graph 参数。这个例子中，我们将实现一个简单的多层感知器（MLP），并使用自定义激活函数进行训练。

import torch
from torch.autograd import Function

# 自定义激活函数ReLU
class MyReLU(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

# 定义模型
class SimpleMLP(torch.nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.linear1 = torch.nn.Linear(2, 4)
        self.linear2 = torch.nn.Linear(4, 1)

    def forward(self, x):
        relu = MyReLU.apply
        x = relu(self.linear1(x))
        x = torch.sigmoid(self.linear2(x))
        return x

# 生成简单的数据集
torch.manual_seed(0)
x_data = torch.tensor([[0.0, 0.0],
                       [0.0, 1.0],
                       [1.0, 0.0],
                       [1.0, 1.0]], requires_grad=True)
y_data = torch.tensor([[0.0], [1.0], [1.0], [0.0]], requires_grad=False)

# 初始化模型和优化器
model = SimpleMLP()
criterion = torch.nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# 训练模型
for epoch in range(1000):
    # 前向传播
    y_pred = model(x_data)
    # 计算损失
    loss = criterion(y_pred, y_data)

    # 反向传播
    optimizer.zero_grad()
    loss.backward(retain_graph=True)
    optimizer.step()

    # 打印损失
    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

# 注册钩子，修改梯度
def hook_fn(grad):
    print('Gradient before modification:', grad)
    return grad * 0.5  # 将梯度值减半

# 选择一个参数进行钩子注册
handle = model.linear1.weight.register_hook(hook_fn)

# 进行一次前向和后向传播以触发钩子
y_pred = model(x_data)
loss = criterion(y_pred, y_data)
optimizer.zero_grad()
loss.backward()

# 移除钩子
handle.remove()

# 再次输出模型参数的梯度以检查修改效果
print('Model parameters after training:')
for param in model.parameters():
    print(param.grad)

# 测试模型
with torch.no_grad():
    print('Testing the model:')
    test_data = torch.tensor([[0.0, 0.0],
                              [0.0, 1.0],
                              [1.0, 0.0],
                              [1.0, 1.0]])
    predictions = model(test_data)
    print(predictions)

代码说明

1. 自定义激活函数 MyReLU：使用 torch.autograd.Function 创建了一个自定义的 ReLU 激活函数，该函数在反向传播中实现了梯度的自定义计算。

2. 定义模型 SimpleMLP：使用两层线性层实现一个简单的多层感知器（MLP）。

3. 数据集：创建了一个简单的 XOR 数据集作为训练和测试数据。

4. 训练循环：在训练过程中，通过调用 loss.backward() 来计算梯度，并使用优化器 optimizer.step() 更新模型参数。

5. 注册钩子：在模型参数上注册一个钩子函数，用于修改反向传播过程中计算得到的梯度。

6. 测试模型：在训练完成后，用测试数据集评估模型的输出。

这个例子展示了如何在 PyTorch 中使用自动微分进行模型训练，并通过钩子和自定义函数增强梯度计算的灵活性。

总结

torch.autograd 提供了灵活而强大的自动微分功能，通过简单地封装张量，并使用 backward() 方法，可以很方便地进行梯度计算。掌握这些技巧可以帮助我们更高效地进行深度学习模型的训练和优化。无论是自定义的损失函数还是特殊的网络结构，自动微分都能帮助我们轻松计算梯度，推动模型的优化和提升。