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\) 是 Hadamard 积。
在多层 LSTM 中，第 \(l\) 层 (\(l \ge 2\)) 的输入 \(x^{(l)}_t\) 是前一层的隐藏状态 \(h^{(l-1)}_t\) 乘以 dropout \(\delta^{(l-1)}_t\)，其中每个 \(\delta^{(l-1)}_t\) 是伯努利随机变量，其概率为
dropout时为 \(0\)。如果指定了
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 ={} & \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已作为输入给出，则输出也将是一个打包序列。当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))
