以下是一个使用 PyTorch 构建 LSTM 模型的简单示例。该模型可以用于序列预测任务，比如时间序列预测或文本生成。 我们将构建一个简单的 LSTM 网络，并训练它学习一个简单的序列任务（例如：根据前几个数字预测下一个数字）。 ### 示例代码 ```python import torch import torch.nn as nn import numpy as np # 定义LSTM模型 class LSTMModel(nn.Module): def __init__(self, input_size=1, hidden_size=50, num_layers=1, output_size=1): super(LSTMModel, self).__init__() self.hidden_size = hidden_size self.num_layers = num_layers # LSTM层 self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) # 全连接层 self.fc = nn.Linear(hidden_size, output_size) def forward(self, x): # 初始化隐藏状态和细胞状态 h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) # 前向传播LSTM out, _ = self.lstm(x, (h0, c0)) # out: tensor of shape (batch_size, seq_length, hidden_size) # 解码最后一个时间步的隐藏状态 out = self.fc(out[:, -1, :]) return out # 超参数 sequence_length = 5 input_size = 1 hidden_size = 50 num_layers = 1 output_size = 1 num_epochs = 100 learning_rate = 0.01 # 创建简单的序列数据：y = x + 1 # 例如输入 [1,2,3,4] 输出 5 def create_data(seq_len=5): X = [] y = [] for i in range(1, 101): X.append(np.array([i + j for j in range(seq_len)])) y.append(i + seq_len) X = np.array(X).reshape(-1, seq_len, 1) y = np.array(y).reshape(-1, 1) return torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32) X, y = create_data(sequence_length) # 实例化模型、损失函数和优化器 model = LSTMModel(input_size, hidden_size, num_layers, output_size) criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # 训练模型 for epoch in range(num_epochs): outputs = model(X) optimizer.zero_grad() loss = criterion(outputs, y) loss.backward() optimizer.step() if (epoch+1) % 10 == 0: print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}') # 测试模型 test_input = torch.tensor([[101, 102, 103, 104, 105]], dtype=torch.float32) predicted = model(test_input) print(f"Predicted next value: {predicted.item():.0f}") ``` --- ### 代码解释： 1. **LSTMModel 类**： - 继承自 `nn.Module`，定义了一个包含 LSTM 层和全连接输出层的神经网络。 - `forward` 方法中，我们传入输入并初始化 LSTM 的初始状态为零。 - 只取最后一个时间步的输出进行预测。 2. **create_data 函数**： - 生成一个简单的递增序列作为训练数据。 - 输入是长度为 `seq_len` 的序列，输出是序列后的一个数字。 3. **模型训练部分**： - 使用均方误差（MSE）作为损失函数。 - 使用 Adam 优化器来更新模型参数。 - 每个 epoch 进行一次前向传播、损失计算和反向传播。 4. **测试模型**： - 输入 `[101, 102, 103, 104, 105]`，模型应输出接近 `106` 的值。 --- ###