LSTMCell¶
- class torchrl.modules.LSTMCell(input_size: int, hidden_size: int, bias: bool = True, device=None, dtype=None)[source]¶
一個長短期記憶 (LSTM) 單元，執行與 nn.LSTMCell 相同的操作，但完全用 Python 程式碼實現。
注意
此類的實現不依賴於 CuDNN，這使得它與
torch.vmap()和torch.compile()相容。示例
>>> import torch >>> from torchrl.modules.tensordict_module.rnn import LSTMCell >>> 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 >>> lstm_cell = LSTMCell(input_size=N_IN, hidden_size=N_OUT, device=device)
# 單次呼叫 >>> x = torch.randn(B, 10, device=device) >>> h0 = torch.zeros(B, 20, device=device) >>> c0 = torch.zeros(B, 20, device=device) >>> with torch.no_grad(): … (h1, c1) = lstm_cell(x, (h0, c0))
# 向量化呼叫 - nn.LSTMCell 不支援 >>> def call_lstm(x, h, c): … h_out, c_out = lstm_cell(x, (h, c)) … return h_out, c_out >>> batched_call = torch.vmap(call_lstm) >>> x = torch.randn(V, B, 10, device=device) >>> h0 = torch.zeros(V, B, 20, device=device) >>> c0 = torch.zeros(V, B, 20, device=device) >>> with torch.no_grad(): … (h1, c1) = batched_call(x, h0, c0)
一個長短期記憶 (LSTM) 單元。
\[\begin{split}\begin{array}{ll} i = \sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\ f = \sigma(W_{if} x + b_{if} + W_{hf} h + b_{hf}) \\ g = \tanh(W_{ig} x + b_{ig} + W_{hg} h + b_{hg}) \\ o = \sigma(W_{io} x + b_{io} + W_{ho} h + b_{ho}) \\ c' = f \odot c + i \odot g \\ h' = o \odot \tanh(c') \\ \end{array}\end{split}\]其中 \(\sigma\) 是 sigmoid 函式，而 \(\odot\) 是 Hadamard 乘積。
- 引數:
input_size – 輸入 x 中預期特徵的數量
hidden_size – 隱藏狀態 h 中的特徵數量
bias – 如果為
False，則該層不使用偏置權重 b_ih 和 b_hh。預設值:True
- 輸入: input, (h_0, c_0)
input 形狀為 (batch, input_size) 或 (input_size): 包含輸入特徵的張量
h_0 形狀為 (batch, hidden_size) 或 (hidden_size): 包含初始隱藏狀態的張量
c_0 形狀為 (batch, hidden_size) 或 (hidden_size): 包含初始單元狀態的張量
如果未提供 (h_0, c_0)，則 h_0 和 c_0 都預設為零。
- 輸出: (h_1, c_1)
h_1 形狀為 (batch, hidden_size) 或 (hidden_size): 包含下一個隱藏狀態的張量
c_1 形狀為 (batch, hidden_size) 或 (hidden_size): 包含下一個單元狀態的張量
- 變數:
weight_ih (torch.Tensor) – 可學習的輸入到隱藏層的權重，形狀為 (4*hidden_size, input_size)
weight_hh (torch.Tensor) – 可學習的隱藏層到隱藏層的權重，形狀為 (4*hidden_size, hidden_size)
bias_ih – 可學習的輸入到隱藏層的偏置，形狀為 (4*hidden_size)
bias_hh – 可學習的隱藏層到隱藏層的偏置，形狀為 (4*hidden_size)
注意
所有權重和偏置均從 \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) 初始化，其中 \(k = \frac{1}{\text{hidden\_size}}\)
在某些 ROCm 裝置上，使用 float16 輸入時，此模組將在反向傳播中使用不同的精度。
示例
>>> rnn = nn.LSTMCell(10, 20) # (input_size, hidden_size) >>> input = torch.randn(2, 3, 10) # (time_steps, batch, input_size) >>> hx = torch.randn(3, 20) # (batch, hidden_size) >>> cx = torch.randn(3, 20) >>> output = [] >>> for i in range(input.size()[0]): ... hx, cx = rnn(input[i], (hx, cx)) ... output.append(hx) >>> output = torch.stack(output, dim=0)
