深入浅出torch.autograd

一、介绍autograd

torch.autograd 模块是 PyTorch 中的自动微分引擎。它支持任意数量的计算图，可以自动执行前向传递、后向传递和计算梯度，同时提供很多有用的函数来处理各种计算图。它与程序中的其它部分紧密结合使用，使得使用 PyTorch 构建神经网络得以简单而高效实现。

二、期望值反向传播

我们简单介绍一下期望值反向传播（REINFORCE)及其PyTorch实现。期望值反向传播是一种用于直接优化policy的算法。在许多强化学习任务中，我们需要优化一个策略，使其能够最大化累积奖励的期望值。在连续动作空间或高维状态空间下，直接优化该策略是非常困难的，因此我们通常使用一些策略梯度方法，它们依据一些目标构建策略的估计，并更新策略以优化该估计。

在期望值反向传播中，我们计算得到目标的期望值，然后对该期望值的对数化为一组权重，分别赋给每一步的回报。然后用一个对策略参数的公式求导。这将给出一组价值函数的梯度，此外，我们还可以使用Torch中的链式法则很容易地获得策略梯度。

import torch

def REINFORCE(policy, episodes, learning_rate):
    optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate) #定义一个优化器
    episode_rewards = []
    for i in range(episodes):
        log_probs = []
        rewards = []

        state = env.reset()
        episode_reward = 0

        while(True):
            action, log_prob = policy.get_action(state)
            new_state, reward, done, _ = env.step(action.item())
            log_probs.append(log_prob)
            rewards.append(reward)
            state = new_state
            episode_reward += reward
            if done:
                break

        episode_rewards.append(episode_reward)

        # 计算期望梯度，更新策略参数
        eps = torch.finfo(torch.float32).eps
        discounts = [np.power(0.99, i) for i in range(len(rewards))]
        discounts = torch.from_numpy(np.array(discounts, dtype=np.float32)).to(device)
        rewards = torch.stack(rewards).to(device)
        log_probs = torch.stack(log_probs).to(device)
        
        R = (rewards * discounts).sum(dim=0)
        R = (R - R.mean()) / (R.std() + eps) # 标准化策略梯度

        policy_loss = (-log_probs * R).sum(dim=0)

        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

    return episode_rewards

三、autograd.Function的使用

本节将详细介绍PyTorch中 autograd.Function 的使用，以及如何使用它们自定义您的操作。在PyTorch中，Variable的每个操作都是如何在计算图中回溯到其他变量，并且每个操作同样可以回溯到其他函数。为了允许用户实现自己的操作，PyTorch为我们提供了一个非常简单和强大的类”。该autograd.Function类被Pytorch用于允许我们定义随时可导的用户自定义运算。

class Exp(Function):

    @staticmethod
    def forward(ctx, i):
        result = i.exp()
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_variables
        return grad_output * result

四、Variable以及在计算图中的应用

在PyTorch中，算子可以添加到计算图中，以实现自动求导。在计算图中，每一个节点都表示一个Tensor，其中一些Tensor节点是输入节点（InputNode），而其他节点是操作节点（FunctionNode）。在轻松编写深度学习模型时，我们很少手动添加节点和边，PyTorch很好地隐藏了这些内容并隐式执行了它。 Variable 是 PyTorch 中图计算的重要概念之一。它是具有梯度的张量，可以直接放入计算图中，并可以通过它的backward()函数产生梯度信号。

Variable是Tensor的一个封装，不同的是Variable有 tensor的一些属性和方法，比如shape，size()等， 同时它有一些附加的属性，比如grad（ Variable的梯度）、requires_grad（是不是需要求 Variable的梯度）、data（保存 Variable的 tensor）等。 我们在需要求 Variable的梯度时才需要调用 backward() 函数.

import torch
from torch.autograd import Variable

x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

w = Variable(torch.Tensor([1.0]), requires_grad=True)

def forward(x):
    return x * w

def loss(x, y):
    y_pred = forward(x)
    return (y_pred - y) * (y_pred - y)

for epoch in range(10):
    for x_val, y_val in zip(x_data, y_data):
        l = loss(x_val, y_val)
        l.backward()
        print("grad: ", x_val, y_val, w.grad.data[0])
        w.data = w.data - 0.01 * w.grad.data

        # Manually zero the gradients after updating weights
        w.grad.data.zero_()

五、tensor.detach()函数的应用

在PyTorch中， tensor.detach()函数是用于获得没有对原始变量的梯度的新张量的，这也就是一个detachdetensor。总之，当您需要获取不需要梯度的张量时， detach() 函数非常有用。在使用GPU时，您必须对张量调用detach()，以便在进行数据移动时清除存储，否则它会导致内存泄漏；因为我们需要计算二阶梯度。 没有从图中分离张量属性。因此， detach() 使您可以在不影响计算图的情况下使用张量。

import torch

tensor = torch.randn(3, requires_grad=True)
print('tensor:', tensor)

# 移动到GPU上
tensor_gpu = tensor.cuda()

# 加点于tensor在GPU上产生的梯度
result_gpu = (tensor_gpu ** 2).sum()
result_gpu.backward()

# TensorFlow，tensor与 tensor_gpu 的梯度不匹配。
# Do a slow transfer back to CPU memory
result = (tensor.detach().cpu() ** 2).sum()
result.backward()

print('tensor_grad:', tensor.grad)
print('tensor_gpu_grad:', tensor_gpu.grad)