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)[源代码]¶
一个 PyTorch 模块，用于执行多层 LSTM 的多个步骤。该模块的行为与
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 随机变量，其值为 \(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})\)（无批处理输入）或 \((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})\) 的张量，包含输入序列中每个元素的初始隐藏状态。如果未提供 (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{批次大小} \\ L ={} & \text{序列长度} \\ D ={} & 2 \text{ 如果 bidirectional=True，否则为 } 1 \\ H_{in} ={} & \text{input\_size} \\ H_{cell} ={} & \text{hidden\_size} \\ H_{out} ={} & \text{proj\_size 如果 } \text{proj\_size}>0 \text{ 否则为 hidden\_size} \\ \end{aligned}\end{split}\]- 输出：output, (h_n, c_n)
output: 形状为 \((L, D * H_{out})\)（无批处理输入）、\((L, N, D * H_{out})\)（当
batch_first=False时）或 \((N, L, D * H_{out})\)（当batch_first=True时）的张量，包含 LSTM 最后一层的输出特征 (h_t)，针对每个 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))
