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]¶
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})\)，当
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{批量大小} \\ L ={} & \text{序列长度} \\ D ={} & 2 \text{ 如果 bidirectional=True 否则为 } 1 \\ H_{in} ={} & \text{输入大小} \\ H_{cell} ={} & \text{细胞大小} \\ H_{out} ={} & \text{投影大小 如果 } \text{投影大小}>0 \text{ 否则为隐藏大小} \\ \end{aligned}\end{split}\]- 输出：output, (h_n, c_n)
output: 包含 LSTM 最后一层在每个时刻 t 的输出特征 (h_t) 的张量，对于非批量输入形状为 \((L, D * H_{out})\)，当
batch_first=False时形状为 \((L, N, D * H_{out})\)，当batch_first=True时形状为 \((N, L, D * H_{out})\)。如果输入是torch.nn.utils.rnn.PackedSequence，则输出也将是 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))
