使用 PyTorch 实现标准 LSTM 神经网络

下面是一篇完整文章,演示如何使用 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}')   
    # 输出标量预测值

逐段总结

  1. 数据处理 :把序列切分为固定长度 seq_length 的输入 X 和目标 Y。

  2. LSTM 定义:一层或多层 LSTM,输出最后时间步的隐藏状态,再用全连接层映射到预测值。

  3. 训练:标准训练流程:清空梯度 → 前向 → 计算损失 → 反向 → 更新参数。

  4. 预测 :用 model.eval() + torch.no_grad(),保证预测时不计算梯度。

  5. 通用性 :这套代码可直接改 input_sizeseq_lengthhidden_sizenum_layersoutput_size 应对不同序列问题。

相关推荐
深度学习lover2 小时前
<项目代码>yolo毛毛虫识别<目标检测>
人工智能·python·yolo·目标检测·计算机视觉·毛毛虫识别
丝斯20112 小时前
AI学习笔记整理(57)——大模型微调相关技术
人工智能·笔记·学习
沃达德软件2 小时前
人脸比对技术助力破案
人工智能·深度学习·神经网络·目标检测·机器学习·生成对抗网络·计算机视觉
副露のmagic2 小时前
Transformer架构
人工智能·深度学习·transformer
softshow10262 小时前
ros2_control
人工智能
cyyt2 小时前
深度学习周报(1.19~1.25)
人工智能·深度学习
2501_948120152 小时前
基于深度学习的遥感图像分类算法研究
人工智能·深度学习·分类
子午2 小时前
【2026计算机毕设】水果识别分类系统~python+深度学习+人工智能+算法模型+TensorFlow
人工智能·python·深度学习
救救孩子把2 小时前
61-机器学习与大模型开发数学教程-5-8 约束非线性优化
人工智能·机器学习