下面是一篇完整文章,演示如何使用 PyTorch 实现 LSTM 神经网络,包括数据准备、模型定义、训练和预测。每行代码都配有注释。
# ====================================
# 导入库
# ====================================
import torch # PyTorch 核心库,负责张量计算和自动求导
import torch.nn as nn # 提供神经网络模块,如 Linear, LSTM 等
import torch.optim as optim # 提供优化器,如 SGD, Adam
from torch.utils.data import DataLoader, TensorDataset
# DataLoader: 用于批量加载数据,支持随机打乱和多线程加速
# TensorDataset: 将特征和标签打包成数据集
# ====================================
# 数据准备
# ====================================
data = [i for i in range(100)]
# 生成简单序列 0-99,用于序列预测示例
seq_length = 5
# 序列长度,即每次用前5个数字预测下一个
X = []
Y = []
for i in range(len(data)-seq_length):
X.append(data[i:i+seq_length]) # 每个输入序列,长度为 seq_length
Y.append(data[i+seq_length]) # 对应的目标值,是序列后的下一个数字
X = torch.tensor(X, dtype=torch.float32).unsqueeze(-1)
# 转成 float32 Tensor
# unsqueeze(-1) 在最后增加一个维度,表示特征维度(input_size=1)
# shape: [样本数, seq_length, 1]
Y = torch.tensor(Y, dtype=torch.float32).unsqueeze(-1)
# 同样处理目标值
# shape: [样本数, 1]
dataset = TensorDataset(X, Y)
# 将特征和标签打包成数据集,方便 DataLoader 使用
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
# 批量加载数据
# batch_size=8: 每次喂入8条序列
# shuffle=True: 每轮训练随机打乱顺序,提高泛化
# ====================================
# 模型定义
# ====================================
class SLTM(nn.Module):
# 定义 LSTM 模型类,继承 nn.Module
def __init__(self, input_size, hidden_size, num_layers, output_size):
super(SLTM, self).__init__()
# 调用父类初始化方法,必须
self.lstm = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True
)
# LSTM 层:
# input_size: 每个时间步输入的特征数
# hidden_size: 隐藏状态维度,影响模型容量
# num_layers: LSTM 堆叠层数
# batch_first=True: 输入形状为 [batch, seq_len, features]
self.fc = nn.Linear(hidden_size, output_size)
# 全连接层,将最后 LSTM 输出映射到最终输出维度
def forward(self, x):
# 前向传播
out, (hn, cn) = self.lstm(x)
# out: LSTM 每个时间步的输出,shape=[batch, seq_len, hidden_size]
# hn: 最后一层隐藏状态,shape=[num_layers, batch, hidden_size]
# cn: 最后一层细胞状态,shape=[num_layers, batch, hidden_size]
out = out[:, -1, :]
# 取最后时间步的输出,作为序列整体的表示
out = self.fc(out)
# 全连接层映射到目标输出
return out
# 返回预测值
# ====================================
# 模型初始化
# ====================================
input_size = 1 # 每个时间步一个特征
hidden_size = 32 # LSTM 隐藏状态神经元数量
num_layers = 1 # LSTM 层数
output_size = 1 # 输出维度
model = SLTM(input_size, hidden_size, num_layers, output_size)
# 实例化模型
criterion = nn.MSELoss()
# 均方误差损失函数,用于回归任务
optimizer = optim.Adam(model.parameters(), lr=0.01)
# Adam 优化器,自动调整学习率,学习率设为0.01
# ====================================
# 训练过程
# ====================================
num_epochs = 200 # 总训练轮数
for epoch in range(num_epochs):
for batch_x, batch_y in dataloader:
optimizer.zero_grad()
# 清空梯度,避免累加
outputs = model(batch_x)
# 前向传播,计算预测值
loss = criterion(outputs, batch_y)
# 计算损失,衡量预测与真实差距
loss.backward()
# 反向传播,计算梯度
optimizer.step()
# 更新参数,梯度下降一步
if (epoch+1) % 20 == 0:
# 每20轮打印一次损失
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
# ====================================
# 预测
# ====================================
model.eval()
# 切换到评估模式,不启用 dropout 或 batchnorm 等训练行为
with torch.no_grad():
# 不需要计算梯度,节省显存
test_seq = torch.tensor([95,96,97,98,99], dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
# 构造预测输入序列
# unsqueeze(0) 添加 batch 维度
# unsqueeze(-1) 添加特征维度
pred = model(test_seq)
# 前向传播得到预测
print(f'预测下一个数字: {pred.item():.4f}')
# 输出标量预测值
逐段总结
-
数据处理 :把序列切分为固定长度
seq_length的输入 X 和目标 Y。 -
LSTM 定义:一层或多层 LSTM,输出最后时间步的隐藏状态,再用全连接层映射到预测值。
-
训练:标准训练流程:清空梯度 → 前向 → 计算损失 → 反向 → 更新参数。
-
预测 :用
model.eval()+torch.no_grad(),保证预测时不计算梯度。 -
通用性 :这套代码可直接改
input_size、seq_length、hidden_size、num_layers、output_size应对不同序列问题。