用pytorch写一个lstm
时间: 2023-04-01 22:03:28 浏览: 159
基于PyTorch框架,使用LSTM网络对股票价格进行预测Python源码
5星 · 资源好评率100%
神经网络模型的实现方式有很多种,其中 LSTM 是一种常用的循环神经网络模型。在 PyTorch 中,可以通过以下代码实现一个简单的 LSTM 模型:
```python
import torch
import torch.nn as nn
class LSTM(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, output_size):
super(LSTM, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
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):
h = torch.zeros(self.num_layers, x.size(), self.hidden_size).to(x.device)
c = torch.zeros(self.num_layers, x.size(), self.hidden_size).to(x.device)
out, _ = self.lstm(x, (h, c))
out = self.fc(out[:, -1, :])
return out
```
其中,`input_size` 表示输入数据的特征维度,`hidden_size` 表示 LSTM 隐藏层的维度,`num_layers` 表示 LSTM 的层数,`output_size` 表示输出数据的维度。在 `forward` 方法中,首先初始化 LSTM 的隐藏状态和细胞状态,然后将输入数据 `x` 传入 LSTM 模型中,得到输出 `out`,最后将 `out` 的最后一个时间步的输出传入全连接层 `fc` 中,得到最终的输出结果。
使用该模型进行训练时,可以通过以下代码实现:
```python
model = LSTM(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):
for i, (inputs, labels) in enumerate(train_loader):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 100 == :
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, num_epochs, i+1, total_step, loss.item()))
```
其中,`train_loader` 是一个 PyTorch 的数据加载器,用于加载训练数据。在每个 epoch 中,对于每个数据批次,首先将输入数据和标签数据转移到指定的设备上,然后将输入数据传入 LSTM 模型中,得到输出结果 `outputs`,计算输出结果和标签数据之间的均方误差损失 `loss`,然后使用反向传播算法更新模型参数。在训练过程中,可以通过打印日志来观察模型的训练情况。
阅读全文