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

相关推荐
墨染天姬4 小时前
【AI】端侧AIBOX可以部署哪些智能体
人工智能
AI成长日志4 小时前
【Agentic RL】1.1 什么是Agentic RL:从传统RL到智能体学习
人工智能·学习·算法
2501_948114244 小时前
2026年大模型API聚合平台技术评测:企业级接入层的治理演进与星链4SAPI架构观察
大数据·人工智能·gpt·架构·claude
小小工匠4 小时前
LLM - awesome-design-md 从 DESIGN.md 到“可对话的设计系统”:用纯文本驱动 AI 生成一致 UI 的新范式
人工智能·ui
黎阳之光5 小时前
黎阳之光:视频孪生领跑者,铸就中国数字科技全球竞争力
大数据·人工智能·算法·安全·数字孪生
小超同学你好5 小时前
面向 LLM 的程序设计 6:Tool Calling 的完整生命周期——从定义、决策、执行到观测回注
人工智能·语言模型
智星云算力5 小时前
本地GPU与租用GPU混合部署:混合算力架构搭建指南
人工智能·架构·gpu算力·智星云·gpu租用
jinanwuhuaguo5 小时前
截止到4月8日,OpenClaw 2026年4月更新深度解读剖析:从“能力回归”到“信任内建”的范式跃迁
android·开发语言·人工智能·深度学习·kotlin
xiaozhazha_5 小时前
效率提升80%:2026年AI CRM与ERP深度集成的架构设计与实现
人工智能
枫叶林FYL5 小时前
【自然语言处理 NLP】7.2.2 安全性评估与Constitutional AI
人工智能·自然语言处理