快捷方式

GRU

class torchrl.modules.GRU(input_size: int, hidden_size: int, num_layers: int = 1, bias: bool, batch_first: bool = True, dropout: float = 0.0, bidirectional: bool = False, device=None, dtype=None)[source]

一個 PyTorch 模組,用於執行多層 GRU 的多個步驟。此模組的行為與 torch.nn.GRU 完全相同,但此實作完全以 Python 編碼。

注意

此類別的實作沒有依賴 CuDNN,因此它與 torch.vmap()torch.compile() 相容。

範例

>>> import torch
>>> from torchrl.modules.tensordict_module.rnn import GRU
>>> device = torch.device("cuda") if torch.cuda.device_count() else torch.device("cpu")
>>> B = 2
>>> T = 4
>>> N_IN = 10
>>> N_OUT = 20
>>> N_LAYERS = 2
>>> V = 4  # vector size
>>> gru = GRU(
...     input_size=N_IN,
...     hidden_size=N_OUT,
...     device=device,
...     num_layers=N_LAYERS,
... )

# 單次呼叫 >>> x = torch.randn(B, T, N_IN, device=device) >>> h0 = torch.zeros(N_LAYERS, B, N_OUT, device=device) >>> with torch.no_grad(): … h1 = gru(x, h0)

# 向量化呼叫 - nn.GRU 無法做到 >>> def call_gru(x, h): … h_out = gru(x, h) … return h_out >>> batched_call = torch.vmap(call_gru) >>> x = torch.randn(V, B, T, 10, device=device) >>> h0 = torch.zeros(V, N_LAYERS, B, N_OUT, device=device) >>> with torch.no_grad(): … h1 = batched_call(x, h0)

__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)

將多層閘控循環單元 (GRU) RNN 應用於輸入序列。對於輸入序列中的每個元素,每一層都會計算以下函數:

\[\begin{split}\begin{array}{ll} r_t = \sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{(t-1)} + b_{hr}) \\ z_t = \sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{(t-1)} + b_{hz}) \\ n_t = \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{(t-1)}+ b_{hn})) \\ h_t = (1 - z_t) \odot n_t + z_t \odot h_{(t-1)} \end{array}\end{split}\]

其中 \(h_t\) 是在時間點 t 的隱藏狀態,\(x_t\) 是在時間點 t 的輸入,\(h_{(t-1)}\) 是在時間點 t-1 的層的隱藏狀態,或是在時間點 0 的初始隱藏狀態,而 \(r_t\)\(z_t\)\(n_t\) 分別是重置閘、更新閘和新閘。\(\sigma\) 是 sigmoid 函數,而 \(\odot\) 是 Hadamard 乘積。

在多層 GRU 中,第 \(l\) 層 (\(l \ge 2\)) 的輸入 \(x^{(l)}_t\) 是前一層的隱藏狀態 \(h^{(l-1)}_t\) 乘以 dropout \(\delta^{(l-1)}_t\),其中每個 \(\delta^{(l-1)}_t\) 是一個 Bernoulli 隨機變數,其機率為 dropout 時為 \(0\)

參數:
  • input_size – 輸入 x 中預期特徵的數量

  • hidden_size – 隱藏狀態 h 中特徵的數量

  • num_layers – 循環層的數量。例如,設定 num_layers=2 表示將兩個 GRU 堆疊在一起以形成一個堆疊 GRU,第二個 GRU 接收第一個 GRU 的輸出並計算最終結果。預設值:1

  • bias – 如果 False,則該層不使用偏置權重 b_ihb_hh。預設值:True

  • batch_first – 如果 True,則輸入和輸出張量會以 (batch, seq, feature) 而不是 (seq, batch, feature) 的形式提供。請注意,這不適用於隱藏狀態或 cell 狀態。有關詳細資訊,請參閱下面的輸入/輸出部分。預設值:False

  • dropout – 如果非零,則在每個 GRU 層(最後一層除外)的輸出上引入一個 Dropout 層,dropout 機率等於 dropout。預設值:0

  • bidirectional – 如果 True,則變成雙向 GRU。預設值:False

輸入:input, h_0
  • input: 形狀為 \((L, H_{in})\) 的張量,用於非批次輸入,\((L, N, H_{in})\) 用於 batch_first=False,或 \((N, L, H_{in})\) 用於 batch_first=True,包含輸入序列的特徵。輸入也可以是一個壓縮的可變長度序列。有關詳細資訊,請參閱 torch.nn.utils.rnn.pack_padded_sequence()torch.nn.utils.rnn.pack_sequence()

  • h_0: 形狀為 \((D * \text{num\_layers}, H_{out})\)\((D * \text{num\_layers}, N, H_{out})\) 的張量,包含輸入序列的初始隱藏狀態。如果未提供,則預設為零。

其中

\[\begin{split}\begin{aligned} N ={} & \text{批次大小} \\ L ={} & \text{序列長度} \\ D ={} & 2 \text{ 如果 bidirectional=True,否則為 } 1 \\ H_{in} ={} & \text{input\_size} \\ H_{out} ={} & \text{hidden\_size} \end{aligned}\end{split}\]
輸出:output, h_n
  • output: 形狀為 \((L, D * H_{out})\) 的張量,用於非批次輸入,\((L, N, D * H_{out})\) 用於 batch_first=False,或 \((N, L, D * H_{out})\) 用於 batch_first=True,包含 GRU 最後一層的輸出特徵 (h_t),對於每個 t。如果給定的輸入是一個 torch.nn.utils.rnn.PackedSequence,則輸出也將是一個壓縮序列。

  • h_n: 形狀為 \((D * \text{num\_layers}, H_{out})\)\((D * \text{num\_layers}, N, H_{out})\) 的張量,包含輸入序列的最終隱藏狀態。

變數:
  • weight_ih_l[k] – 第 \(\text{k}^{th}\) 層的可學習輸入-隱藏權重 (W_ir|W_iz|W_in),對於 k = 0,形狀為 (3*hidden_size, input_size)。否則,形狀為 (3*hidden_size, num_directions * hidden_size)

  • weight_hh_l[k] – 第 \(\text{k}^{th}\) 層的可學習隱藏-隱藏權重 (W_hr|W_hz|W_hn),形狀為 (3*hidden_size, hidden_size)

  • bias_ih_l[k] – 第 \(\text{k}^{th}\) 層的可學習輸入-隱藏偏置 (b_ir|b_iz|b_in),形狀為 (3*hidden_size)

  • bias_hh_l[k] – 第 \(\text{k}^{th}\) 層的可學習隱藏-隱藏偏置 (b_hr|b_hz|b_hn),形狀為 (3*hidden_size)

注意

所有權重和偏置都從 \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) 初始化,其中 \(k = \frac{1}{\text{hidden\_size}}\)

注意

對於雙向 GRU,前向和後向分別是方向 0 和 1。當 batch_first=False 時,分割輸出層的範例:output.view(seq_len, batch, num_directions, hidden_size)

注意

對於非批次輸入,batch_first 參數將被忽略。

注意

新閘 \(n_t\) 的計算與原始論文和其他框架略有不同。在原始實現中,\(r_t\) 和先前的隱藏狀態 \(h_{(t-1)}\) 之間的 Hadamard 乘積 \((\odot)\) 是在與權重矩陣 W 相乘並加上偏置之前完成的。

\[\begin{aligned} n_t = \tanh(W_{in} x_t + b_{in} + W_{hn} ( r_t \odot h_{(t-1)} ) + b_{hn}) \end{aligned}\]

這與 PyTorch 的實現相反,後者是在 \(W_{hn} h_{(t-1)}\) 之後完成的。

\[\begin{aligned} n_t = \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{(t-1)}+ b_{hn})) \end{aligned}\]

為了效率起見,此實現故意有所不同。

範例

>>> rnn = nn.GRU(10, 20, 2)
>>> input = torch.randn(5, 3, 10)
>>> h0 = torch.randn(2, 3, 20)
>>> output, hn = rnn(input, h0)
forward(input, hx=None)[原始碼]

定義每次呼叫時執行的計算。

應被所有子類別覆寫。

注意

雖然 forward 傳遞的步驟需要在這個函式中定義,但應該在之後呼叫 Module 實例,而不是直接呼叫這個函式,因為前者會處理已註冊的 hooks,而後者會默默地忽略它們。

文件

存取 PyTorch 的完整開發者文件

檢視文件

教學

獲取初學者和高級開發者的深度教學

檢視教學

資源

尋找開發資源並獲得問題解答

檢視資源