Pytorch 的 LSTM 模型的简单示例_lstm代码

1. 代码

完整的源代码：

import torch
from torch import nn

# 定义一个LSTM模型
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):
        # 初始化隐藏状态h0, c0为全0向量
        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)

        # 将输入x和隐藏状态(h0, c0)传入LSTM网络
        out, _ = self.lstm(x, (h0, c0))
        # 取最后一个时间步的输出作为LSTM网络的输出
        out = self.fc(out[:, -1, :])
        return out

# 定义LSTM超参数
input_size = 10   # 输入特征维度
hidden_size = 32  # 隐藏单元数量
num_layers = 2    # LSTM层数
output_size = 2   # 输出类别数量

# 构建一个随机输入x和对应标签y
x = torch.randn(64, 5, 10)  # [batch_size, sequence_length, input_size]
y = torch.randint(0, 2, (64,))  # 二分类任务，标签为0或1

# 创建LSTM模型，并将输入x传入模型计算预测输出
lstm = LSTM(input_size, hidden_size, num_layers, output_size)
pred = lstm(x)  # [batch_size, output_size]

# 定义损失函数和优化器，并进行模型训练
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(lstm.parameters(), lr=1e-3)
num_epochs = 100

for epoch in range(num_epochs):
    # 前向传播计算损失函数值
    pred = lstm(x)  # 在每个epoch中重新计算预测输出
    loss = criterion(pred.squeeze(), y)

    # 反向传播更新模型参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 输出每个epoch的训练损失
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

2. 模型结构分析

# 定义一个LSTM模型
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):
        # 初始化隐藏状态h0, c0为全0向量
        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)

        # 将输入x和隐藏状态(h0, c0)传入LSTM网络
        out, _ = self.lstm(x, (h0, c0))
        # 取最后一个时间步的输出作为LSTM网络的输出
        out = self.fc(out[:, -1, :])
        return out
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

上述代码定义了一个LSTM类，这个类可以用于完成一个基于LSTM的序列模型的搭建。

在初始化函数中，输入的参数分别是输入数据的特征维度（input_size），隐藏层的大小（hidden_size），LSTM层数（num_layers）以及输出数据的维度（output_size）。这里使用batch_first=True表示输入数据的第一个维度是batch size，第二个维度是时间步长和特征维度。

在forward函数中，首先初始化了LSTM网络的隐藏状态为全0向量，并且将其移动到与输入数据相同的设备上。然后调用了nn.LSTM函数进行前向传播操作，并且通过fc层将最后一个时间步的输出映射为输出的数据，最后进行了返回。

3. 代码详解

        # 将输入x和隐藏状态(h0, c0)传入LSTM网络
        out, _ = self.lstm(x, (h0, c0))
1
2

这行代码是利用 PyTorch 自带的 LSTM 模块处理输入张量 x（形状为 [batch_size, sequence_length, input_size]）并得到 LSTM 层的输出 out 和最终状态。其中，h0 是 LSTM 层的初始隐藏状态，c0 是 LSTM 层的初始细胞状态。

在代码中，调用了 self.lstm(x, (h0, c0)) 函数，该函数的返回值有两个：第一个返回值是 LSTM 层的输出 out，其包含了所有时间步上的隐状态；第二个返回值是一个元组，包含了最后一个时间步的隐藏状态和细胞状态，但我们用“_”丢弃了它。

因为对于许多深度学习任务来说，只需要输出序列的最后一个时间步的隐藏状态，而不需要每个时间步上的隐藏状态。因此，这里我们只保留 LSTM 层的输出 out，而忽略了 LSTM 层最后时间步的状态。

最后，out 的形状为 [batch_size, sequence_length, hidden_size]，其中 hidden_size 是 LSTM 层输出的隐藏状态的维度大小。

x = torch.randn(64, 5, 10)
1

这行代码创建了一个形状为 (64, 5, 10) 的张量 x，它包含 64 个样本，每个样本具有 5 个特征维度和 10 个时间步。该张量的值是由均值为 0，标准差为 1 的正态分布随机生成的。

torch.randn() 是 PyTorch 中生成服从标准正态分布的随机数的函数。它的输入是张量的形状，输出是符合正态分布的张量。在本例中，形状为 (64, 5, 10) 表示该张量包含 64 个样本，每个样本包含 5 个特征维度和 10 个时间步，每个元素都是服从标准正态分布的随机数。这种方式生成的随机数可以用于初始化模型参数、生成噪音数据等许多深度学习应用场景。

y = torch.randint(0, 2, (64,))  # 二分类任务，标签为0或1
1

y = torch.randint(0, 2, (64,)) 是使用 PyTorch 库中的 randint() 函数来生成一个64个元素的张量 y，张量的每个元素都是从区间 [0, 2) 中随机生成的整数。

具体而言，torch.randint() 函数包含三个参数，分别是 low、high 和 size。其中，low 和 high 分别表示随机生成整数的区间为 [low, high)，而 size 参数指定了生成的张量的形状。

在上述代码中，size=(64,) 表示生成的张量 y 的形状为 64x1，即一个包含 64 个元素的一维张量，并且每个元素的值都在 [0, 2) 中随机生成。这种形式的张量通常用于分类问题中的标签向量。在该任务中，一个标签通常由一个整数表示，因此可以采用使用 randint() 函数生成一个长度为标签类别数的一维张量，其每个元素的取值为 0 或 1，表示对应类别是否被选中。

# 创建LSTM模型，并将输入x传入模型计算预测输出
lstm = LSTM(input_size, hidden_size, num_layers, output_size)
pred = lstm(x)  # [batch_size, output_size]
1
2
3

通过定义的LSTM类创建了一个LSTM模型，并将输入x传入模型进行前向计算，得到了一个预测输出pred，其形状为[64, output_size]，其中output_size是在LSTM初始化函数中指定的输出数据的维度。

这段代码演示了如何使用已经构建好的代码搭建并训练一个基于LSTM的序列模型，并且展示了其中的一些关键步骤，包括数据输入、模型创建以及前向计算。