当前位置：首页 > article >正文

pytorch LSTM 结构详解

article 2025/9/1 12:28:54

最近项目用到了LSTM ，但是对LSTM 的输入输出不是很理解，对此，我详细查找了lstm 的资料

import torch.nn as nnclass LSTMModel(nn.Module):def __init__(self, input_size=1, hidden_size=50, num_layers=2):super(LSTMModel, self).__init__()self.hidden_size = hidden_sizeself.num_layers = num_layersself.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)self.fc = nn.Linear(hidden_size, 1)  # 1 表示预测输出变量为1def forward(self, x):h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size)c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size)out, _ = self.lstm(x, (h0, c0))out = self.fc(out[:, -1, :])return out # out 形状为（batch_size,1)

input_size=1：输入特征的维度，适用于单变量时间序列。
hidden_size=50：LSTM 隐藏层的维度，决定了模型的记忆能力。
num_layers=2：堆叠的 LSTM 层数，增加层数可以提升模型的表达能力。
batch_first=True：指定输入和输出的张量形状为 (batch_size, seq_len, input_size)。
self.fc：一个全连接层，将 LSTM 的输出映射到最终的预测值。
batch_size 表示批次、seq_len 表示窗口大小、input_size 表示输入尺寸，单变量输入为1 ，多变量要基于个数变化

初始化隐藏状态和细胞状态：
- h0 和 c0 分别表示初始的隐藏状态和细胞状态，形状为 (num_layers, batch_size, hidden_size)。
- 在每次前向传播时，初始化为零张量。
LSTM 层处理：
- self.lstm(x, (h0, c0))：将输入 x 和初始状态传入 LSTM 层，输出 out 和新的状态。
- out 的形状为 (batch_size, seq_len, hidden_size)，包含了每个时间步的输出。
全连接层映射：
- out[:, -1, :]：提取序列中最后一个时间步的输出。
- self.fc(...)：将提取的输出通过全连接层，得到最终的预测结果。