深入淺出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)