LSTM¶
- class torchrl.modules.LSTM(input_size: int, hidden_size: int, num_layers: int = 1, batch_first: bool = True, bias: bool = True, dropout: float = 0.0, bidirectional: float = False, proj_size: int = 0, device=None, dtype=None)[source]¶
一个用于执行多层 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\) 是哈达玛积。
在多层 LSTM 中，第 \(l\) 层 (\(l \ge 2\)) 的输入 \(x^{(l)}_t\) 是前一层隐藏状态 \(h^{(l-1)}_t\) 乘以 dropout \(\delta^{(l-1)}_t\)，其中每个 \(\delta^{(l-1)}_t\) 都是一个伯努利随机变量，其概率为 \(0\)，概率为
dropout。如果指定了
proj_size > 0，则将使用带投影的 LSTM。这将以下列方式更改 LSTM 单元格。首先，\(h_t\) 的维度将从hidden_size更改为proj_size（\(W_{hi}\) 的维度将相应更改）。其次，每一层的输出隐藏状态将乘以一个可学习的投影矩阵：\(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 堆叠在一起形成一个“堆叠 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，则将使用具有相应大小投影的 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})\)，包含输入序列的特征。输入也可以是打包的可变长度序列。有关详细信息，请参阅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 ={} & 批次大小 \\ L ={} & 序列长度 \\ D ={} & 如果 bidirectional=True 则为 2，否则为 1 \\ H_{in} ={} & input\_size \\ H_{cell} ={} & hidden\_size \\ H_{out} ={} & 如果 proj\_size>0 则为 proj\_size，否则为 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 最后一层每个 t 的输出特征 (h_t)。如果已将torch.nn.utils.rnn.PackedSequence作为输入给出，则输出也将是打包的序列。当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)，对于 k = 0，形状为 (4*hidden_size, input_size)。否则，形状为 (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))
