LSTM¶
- class torchrl.modules.LSTM(input_size: int, hidden_size: int, num_layers: int = 1, batch_first: bool, bias: bool = True, dropout: float = 0.0, bidirectional: float = False, proj_size: int = 0, device=None, dtype=None)[來源]¶
一個用於執行多層 LSTM 多個步驟的 PyTorch 模組。該模組的行為與
torch.nn.LSTM完全相同，但此實現完全以 Python 編碼。注意
這個類別的實作不依賴 CuDNN，因此它與
torch.vmap()和torch.compile()相容。範例
>>> import torch >>> from torchrl.modules.tensordict_module.rnn import LSTM
>>> 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 >>> lstm = LSTM( ... 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) >>> c0 = torch.zeros(N_LAYERS, B, N_OUT, device=device) >>> with torch.no_grad(): … h1, c1 = lstm(x, (h0, c0))
# 向量化呼叫 - nn.LSTM 無法做到 >>> def call_lstm(x, h, c): … h_out, c_out = lstm(x, (h, c)) … return h_out, c_out >>> batched_call = torch.vmap(call_lstm) >>> x = torch.randn(V, B, T, 10, device=device) >>> h0 = torch.zeros(V, N_LAYERS, B, N_OUT, device=device) >>> c0 = torch.zeros(V, N_LAYERS, B, N_OUT, device=device) >>> with torch.no_grad(): … h1, c1 = batched_call(x, h0, c0)
__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,proj_size=0,device=None,dtype=None)
將多層長短期記憶 (LSTM) RNN 應用於輸入序列。對於輸入序列中的每個元素，每一層計算以下函數
\[\begin{split}\begin{array}{ll} \\ i_t = \sigma(W_{ii} x_t + b_{ii} + W_{hi} h_{t-1} + b_{hi}) \\ f_t = \sigma(W_{if} x_t + b_{if} + W_{hf} h_{t-1} + b_{hf}) \\ g_t = \tanh(W_{ig} x_t + b_{ig} + W_{hg} h_{t-1} + b_{hg}) \\ o_t = \sigma(W_{io} x_t + b_{io} + W_{ho} h_{t-1} + b_{ho}) \\ c_t = f_t \odot c_{t-1} + i_t \odot g_t \\ h_t = o_t \odot \tanh(c_t) \\ \end{array}\end{split}\]其中 \(h_t\) 是時間 t 的隱藏狀態，\(c_t\) 是時間 t 的細胞狀態，\(x_t\) 是時間 t 的輸入，\(h_{t-1}\) 是時間 t-1 的層的隱藏狀態，或是時間 0 的初始隱藏狀態，而 \(i_t\)、\(f_t\)、\(g_t\)、\(o_t\) 分別是輸入閘、遺忘閘、細胞閘和輸出閘。\(\sigma\) 是 sigmoid 函數，而 \(\odot\) 是 Hadamard 乘積。
在多層 LSTM 中，第 \(l\) 層（\(l \ge 2\)）的輸入 \(x^{(l)}_t\) 是前一層的隱藏狀態 \(h^{(l-1)}_t\) 乘以 dropout \(\delta^{(l-1)}_t\)，其中每個 \(\delta^{(l-1)}_t\) 都是一個 Bernoulli 隨機變數，其機率為
dropout時為 \(0\)。如果指定了
proj_size > 0，將使用帶有 projection 的 LSTM。這會以以下方式更改 LSTM 細胞。首先，\(h_t\) 的維度將從hidden_size變更為proj_size（\(W_{hi}\) 的維度也會相應變更）。其次，每一層的輸出隱藏狀態將乘以可學習的 projection 矩陣：\(h_t = W_{hr}h_t\)。請注意，由於這個結果，LSTM 網路的輸出也將具有不同的形狀。有關所有變數的確切維度，請參閱下面的輸入/輸出部分。您可以在 https://arxiv.org/abs/1402.1128 中找到更多詳細資訊。- 參數:
input_size – 輸入 x 中預期特徵的數量
hidden_size – 隱藏狀態 h 中特徵的數量
num_layers – 遞迴層的數量。例如，設定
num_layers=2表示將兩個 LSTM 堆疊在一起以形成一個 stacked LSTM，第二個 LSTM 接收第一個 LSTM 的輸出並計算最終結果。預設值：1bias – 如果
False，則該層不使用偏差權重 b_ih 和 b_hh。預設值：Truebatch_first – 如果
True，則輸入和輸出張量以 (batch, seq, feature) 而非 (seq, batch, feature) 的形式提供。 請注意，這不適用於隱藏或細胞狀態。有關詳細資訊，請參閱下面的輸入/輸出部分。預設值：Falsedropout – 如果非零，則在每個 LSTM 層（除了最後一層）的輸出上引入一個 Dropout 層，dropout 機率等於
dropout。預設值：0bidirectional – 如果
True，則變成雙向 LSTM。預設值：Falseproj_size – 如果
> 0，將使用具有相應大小 projection 的 LSTM。預設值：0
- 輸入：input, (h_0, c_0)
input：對於未批次處理的輸入，形狀為 \((L, H_{in})\) 的張量，當
batch_first=False時，形狀為 \((L, N, H_{in})\) 的張量，或者當batch_first=True時，形狀為 \((N, L, H_{in})\) 的張量，包含輸入序列的特徵。輸入也可以是 packed 可變長度序列。 有關詳細資訊，請參閱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})\) 的張量，包含輸入序列中每個元素的初始隱藏狀態。如果未提供 (h_0, c_0)，則預設為零。
c_0：對於未批次處理的輸入，形狀為 \((D * \text{num\_layers}, H_{cell})\) 的張量，或者形狀為 \((D * \text{num\_layers}, N, H_{cell})\) 的張量，包含輸入序列中每個元素的初始細胞狀態。如果未提供 (h_0, c_0)，則預設為零。
其中
\[\begin{split}\begin{aligned} N ={} & \text{batch size} \\ L ={} & \text{sequence length} \\ D ={} & 2 \text{ if bidirectional=True otherwise } 1 \\ H_{in} ={} & \text{input\_size} \\ H_{cell} ={} & \text{hidden\_size} \\ H_{out} ={} & \text{proj\_size if } \text{proj\_size}>0 \text{ otherwise hidden\_size} \\ \end{aligned}\end{split}\]- 輸出：output, (h_n, c_n)
output：形狀為 \((L, D * H_{out})\) 的張量，用於非批次輸入；當
batch_first=False時，形狀為 \((L, N, D * H_{out})\)；當batch_first=True時，形狀為 \((N, L, D * H_{out})\)，包含 LSTM 最後一層的輸出特徵 (h_t)，針對每個 t。如果輸入為torch.nn.utils.rnn.PackedSequence，則輸出也將是一個 packed sequence。當bidirectional=True時，output 將包含序列中每個時間步的前向和反向隱藏狀態的串聯。h_n：形狀為 \((D * \text{num\_layers}, H_{out})\) 的張量，用於非批次輸入；或形狀為 \((D * \text{num\_layers}, N, H_{out})\) 的張量，包含序列中每個元素的最終隱藏狀態。當
bidirectional=True時，h_n 將包含最終前向和反向隱藏狀態的串聯。c_n：形狀為 \((D * \text{num\_layers}, H_{cell})\) 的張量，用於非批次輸入；或形狀為 \((D * \text{num\_layers}, N, H_{cell})\) 的張量，包含序列中每個元素的最終記憶胞狀態。當
bidirectional=True時，c_n 將包含最終前向和反向記憶胞狀態的串聯。
- 變數:
weight_ih_l[k] – 第 \(\text{k}^{th}\) 層的可學習輸入到隱藏層的權重 (W_ii|W_if|W_ig|W_io)，形狀為 (4*hidden_size, input_size)，當 k = 0 時。否則，形狀為 (4*hidden_size, num_directions * hidden_size)。如果指定了
proj_size > 0，則當 k > 0 時，形狀將為 (4*hidden_size, num_directions * proj_size)。weight_hh_l[k] – 第 \(\text{k}^{th}\) 層的可學習隱藏層到隱藏層的權重 (W_hi|W_hf|W_hg|W_ho)，形狀為 (4*hidden_size, hidden_size)。如果指定了
proj_size > 0，則形狀將為 (4*hidden_size, proj_size)。bias_ih_l[k] – 第 \(\text{k}^{th}\) 層的可學習輸入到隱藏層的偏差 (b_ii|b_if|b_ig|b_io)，形狀為 (4*hidden_size)。
bias_hh_l[k] – 第 \(\text{k}^{th}\) 層的可學習隱藏層到隱藏層的偏差 (b_hi|b_hf|b_hg|b_ho)，形狀為 (4*hidden_size)。
weight_hr_l[k] – 第 \(\text{k}^{th}\) 層的可學習投影權重，形狀為 (proj_size, hidden_size)。僅當指定了
proj_size > 0時存在。weight_ih_l[k]_reverse – 與反向方向的 weight_ih_l[k] 類似。僅當
bidirectional=True時存在。weight_hh_l[k]_reverse – 與反向方向的 weight_hh_l[k] 類似。僅當
bidirectional=True時存在。bias_ih_l[k]_reverse – 與反向方向的 bias_ih_l[k] 類似。僅當
bidirectional=True時存在。bias_hh_l[k]_reverse – 與反向方向的 bias_hh_l[k] 類似。僅當
bidirectional=True時存在。weight_hr_l[k]_reverse – 與反向方向的 weight_hr_l[k] 類似。僅當指定了
bidirectional=True和proj_size > 0時存在。
注意
所有權重和偏差都從 \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) 初始化，其中 \(k = \frac{1}{\text{hidden\_size}}\)
注意
對於雙向 LSTM，前向和後向分別為方向 0 和 1。 當
batch_first=False時，拆分輸出層的範例：output.view(seq_len, batch, num_directions, hidden_size)。注意
對於雙向 LSTM，h_n 不等於 output 的最後一個元素；前者包含最終的前向和反向隱藏狀態，而後者包含最終的前向隱藏狀態和初始的反向隱藏狀態。
注意
對於非批次輸入，
batch_first參數會被忽略。注意
proj_size應小於hidden_size。範例
>>> rnn = nn.LSTM(10, 20, 2) >>> input = torch.randn(5, 3, 10) >>> h0 = torch.randn(2, 3, 20) >>> c0 = torch.randn(2, 3, 20) >>> output, (hn, cn) = rnn(input, (h0, c0))
