使用 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 应对不同序列问题。

相关推荐
江边风声13 分钟前
从薄板到厚板、从吸盘到夹板边——坤鹏伯爵的取放技术体系是怎样覆盖全制程的
人工智能·科技·自动化·制造·pcb工艺
商业模式源码开发17 分钟前
小米新车未发布即遭 AI 谣言攻击:黑色 GEO 的运作原理与企业正规防御方案
大数据·人工智能·ai·geo
edtoplort30 分钟前
A股最大IPO长鑫科技295亿:从年亏163亿到日赚3亿
大数据·人工智能·科技
道影子37 分钟前
《道德经》031兵者不祥,胜以丧礼处之
人工智能·深度学习·算法
MGS浪疯41 分钟前
我用 WorkBuddy 做了一个能记住前文的 AI 小说编辑器
人工智能·编辑器·腾讯云·腾讯云ai代码助手·workbuddy开发者分享季·codebuddy开发者分享季
IT_陈寒42 分钟前
Python的线程池把我CPU跑满了,原来少传了个参数
前端·人工智能·后端
Microvision维视智造1 小时前
买视觉系统踩过的坑,都是选型时没问对的问题,真实案例手把手教你选择视觉合作伙伴!
人工智能·计算机视觉·视觉检测·机器视觉
BSD_CGQ1 小时前
FSR压力传感器的MCU采集方案与校准算法实现
人工智能·计算机视觉·目标跟踪·adc·压力传感器·fsr·深圳源头厂家
老猿AI洞察1 小时前
阿里Qwen-Image-3.0发布:4.5K token长输入+10px小字渲染,AI生图终于能从“好看“走向“好用“了
人工智能·阿里云
愚公搬代码1 小时前
测试测试测试测试测试测试
人工智能