GRUCell¶
- class torchrl.modules.GRUCell(input_size: int, hidden_size: int, bias: bool = True, device=None, dtype=None)[原始碼]¶
一個閘控循環單元 (GRU) cell，其執行與 nn.LSTMCell 相同的操作，但完全以 Python 編碼。
注意
這個類別的實作不依賴 CuDNN，這使其與
torch.vmap()和torch.compile()相容。範例
>>> import torch >>> from torchrl.modules.tensordict_module.rnn import GRUCell >>> device = torch.device("cuda") if torch.cuda.device_count() else torch.device("cpu") >>> B = 2 >>> N_IN = 10 >>> N_OUT = 20 >>> V = 4 # vector size >>> gru_cell = GRUCell(input_size=N_IN, hidden_size=N_OUT, device=device)
# 單一呼叫 >>> x = torch.randn(B, 10, device=device) >>> h0 = torch.zeros(B, 20, device=device) >>> with torch.no_grad(): … h1 = gru_cell(x, h0)
# 向量化呼叫 - nn.GRUCell 無法使用 >>> def call_gru(x, h): … h_out = gru_cell(x, h) … return h_out >>> batched_call = torch.vmap(call_gru) >>> x = torch.randn(V, B, 10, device=device) >>> h0 = torch.zeros(V, B, 20, device=device) >>> with torch.no_grad(): … h1 = batched_call(x, h0)
一個閘控循環單元 (GRU) cell。
\[\begin{split}\begin{array}{ll} r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\ z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\ n = \tanh(W_{in} x + b_{in} + r \odot (W_{hn} h + b_{hn})) \\ h' = (1 - z) \odot n + z \odot h \end{array}\end{split}\]其中 \(\sigma\) 是 sigmoid 函數，而 \(\odot\) 是 Hadamard 乘積。
- 參數:
input_size – 輸入 x 中預期特徵的數量。
hidden_size – 隱藏狀態 h 中的特徵數量。
bias – 如果
False，則該層不使用偏差權重 b_ih 和 b_hh。預設值：True
- 輸入：input, hidden
input：包含輸入特徵的張量 (tensor)。
hidden：包含批次中每個元素的初始隱藏狀態的張量 (tensor)。如果未提供，則預設為零。
- 輸出：h’
h’：包含批次中每個元素的下一個隱藏狀態的張量 (tensor)。
- 形狀
input: 包含輸入特徵的 \((N, H_{in})\) 或 \((H_{in})\) 張量 (tensor)，其中 \(H_{in}\) = input_size。
hidden: 包含初始隱藏狀態的 \((N, H_{out})\) 或 \((H_{out})\) 張量 (tensor)，其中 \(H_{out}\) = hidden_size。如果未提供，則預設為零。
output: 包含下一個隱藏狀態的 \((N, H_{out})\) 或 \((H_{out})\) 張量 (tensor)。
- 變數:
weight_ih (torch.Tensor) – 可學習的輸入-隱藏權重，形狀為 (3*hidden_size, input_size)
weight_hh (torch.Tensor) – 可學習的隱藏-隱藏權重，形狀為 (3*hidden_size, hidden_size)
bias_ih – 可學習的輸入-隱藏偏差，形狀為 (3*hidden_size)
bias_hh – 可學習的隱藏-隱藏偏差，形狀為 (3*hidden_size)
注意
所有權重和偏差都從 \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) 初始化，其中 \(k = \frac{1}{\text{hidden\_size}}\)
在某些 ROCm 裝置上，當使用 float16 輸入時，此模組將使用不同精度進行反向傳播。
範例
>>> rnn = nn.GRUCell(10, 20) >>> input = torch.randn(6, 3, 10) >>> hx = torch.randn(3, 20) >>> output = [] >>> for i in range(6): ... hx = rnn(input[i], hx) ... output.append(hx)
