1. PyTorch 中的 LSTM

PyTorch 已经封装好了 LSTM：

import torch
import torch.nn as nn

lstm = nn.LSTM(
    input_size=100,
    hidden_size=128,
    num_layers=1,
    batch_first=True
)

这里最重要的参数有 4 个。

1. input_size

表示每个时间步输入特征的维度。

比如：

文本任务中：
每个词经过 embedding 后是 100 维
那么 input_size = 100
时间序列任务中：
每个时间点有 5 个特征
那么 input_size = 5
音频任务中：
每一帧声学特征是 74 维
那么 input_size = 74

2. hidden_size

表示 LSTM 隐藏状态的维度。

例如：

hidden_size=128

表示每个时间步输出一个 128 维的隐藏表示。

如果输入是：

x.shape = [32, 20, 100]

其中：

32：batch_size
20：seq_len
100：input_size

那么输出是：

output.shape = [32, 20, 128]

3. num_layers

表示 LSTM 堆叠几层。

num_layers=1

表示单层 LSTM。

num_layers=2

表示两层 LSTM。

一般建议：

初学和小数据集：1 层
常规任务：1 到 2 层
不要一开始就堆很多层，容易过拟合

4. batch_first

这个参数非常重要。

如果：

batch_first=True

输入格式是：

[batch_size, seq_len, input_size]

如果：

batch_first=False

输入格式是：

[seq_len, batch_size, input_size]

我们一般推荐：

batch_first=True

因为它更符合普通深度学习任务的数据习惯。

2. 最小 LSTM 示例

import torch
import torch.nn as nn

batch_size = 32
seq_len = 20
input_size = 100
hidden_size = 128

x = torch.randn(batch_size, seq_len, input_size)

lstm = nn.LSTM(
    input_size=input_size,
    hidden_size=hidden_size,
    num_layers=1,
    batch_first=True
)

output, (h_n, c_n) = lstm(x)

print("output shape:", output.shape)
print("h_n shape:", h_n.shape)
print("c_n shape:", c_n.shape)

输出结果是：

output shape: torch.Size([32, 20, 128])
h_n shape: torch.Size([1, 32, 128])
c_n shape: torch.Size([1, 32, 128])

3. output、h_n、c_n 分别是什么？

这是 PyTorch LSTM 最容易混的地方。

1. output

output.shape = [batch_size, seq_len, hidden_size]

它表示：

每一个时间步的隐藏状态。

也就是：

h1, h2, h3, ..., hT

如果：

output.shape = [32, 20, 128]

表示：

32 个样本
每个样本有 20 个时间步
每个时间步输出 128 维表示

2. h_n

h_n.shape = [num_layers, batch_size, hidden_size]

它表示：

最后一个时间步的隐藏状态。

如果只有一层 LSTM：

h_n.shape = [1, 32, 128]

那么：

h_n[-1].shape = [32, 128]

这个通常可以作为整个序列的表示。

3. c_n

c_n.shape = [num_layers, batch_size, hidden_size]

它表示：

最后一个时间步的细胞状态，也就是长期记忆。

一般分类任务中，我们更常用：

h_n[-1]

而不是 c_n。

4. output 和 h_n 的区别

对于单层单向 LSTM：

output[:, -1, :]

和：

h_n[-1]

通常是一样的。

也就是：

last_output = output[:, -1, :]
last_hidden = h_n[-1]

二者形状都是：

[batch_size, hidden_size]

但是注意：

在多层、双向、变长序列情况下，二者含义会有细微差别。

初学阶段你先记住：

做序列分类时，可以优先使用 h_n[-1]。

5. 用 LSTM 做文本分类

文本分类流程一般是：

输入词 id
↓
Embedding
↓
LSTM
↓
取最后隐藏状态
↓
全连接层
↓
输出类别 logits

模型代码

import torch
import torch.nn as nn


class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_classes):
        super().__init__()

        self.embedding = nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=embed_dim
        )

        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True
        )

        self.classifier = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # x: [batch_size, seq_len]

        embedded = self.embedding(x)
        # embedded: [batch_size, seq_len, embed_dim]

        output, (h_n, c_n) = self.lstm(embedded)
        # output: [batch_size, seq_len, hidden_size]
        # h_n: [1, batch_size, hidden_size]

        final_hidden = h_n[-1]
        # final_hidden: [batch_size, hidden_size]

        logits = self.classifier(final_hidden)
        # logits: [batch_size, num_classes]

        return logits

6. 跑一个假数据测试

vocab_size = 5000
embed_dim = 100
hidden_size = 128
num_classes = 2

model = LSTMClassifier(
    vocab_size=vocab_size,
    embed_dim=embed_dim,
    hidden_size=hidden_size,
    num_classes=num_classes
)

x = torch.randint(0, vocab_size, (32, 20))
# x: [batch_size, seq_len]

logits = model(x)

print(logits.shape)

输出：

torch.Size([32, 2])

含义是：

32 个样本
每个样本输出 2 个类别分数

对于二分类任务，两个类别可以是：

negative
positive

7. 损失函数怎么写？

如果是多分类或二分类的类别索引形式，使用：

criterion = nn.CrossEntropyLoss()

模型输出：

logits.shape = [batch_size, num_classes]

标签：

labels.shape = [batch_size]

例如：

labels = torch.randint(0, 2, (32,))

训练一步：

criterion = nn.CrossEntropyLoss()

logits = model(x)
loss = criterion(logits, labels)

print(loss)

注意：

CrossEntropyLoss 接收的是 logits，不需要提前 softmax。

错误写法：

probs = torch.softmax(logits, dim=-1)
loss = criterion(probs, labels)

正确写法：

loss = criterion(logits, labels)

8. 优化器怎么写？

常用 Adam：

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3
)

训练步骤：

optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, labels)
loss.backward()
optimizer.step()

完整一轮流程：

model.train()

for x, labels in train_loader:
    optimizer.zero_grad()

    logits = model(x)
    loss = criterion(logits, labels)

    loss.backward()
    optimizer.step()

9. 加 Dropout 的 LSTM 分类器

为了防止过拟合，可以加入 dropout。

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_classes, dropout=0.5):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embed_dim)

        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True
        )

        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        embedded = self.embedding(x)

        output, (h_n, c_n) = self.lstm(embedded)

        final_hidden = h_n[-1]
        final_hidden = self.dropout(final_hidden)

        logits = self.classifier(final_hidden)

        return logits

注意：

nn.LSTM 自带的 dropout 参数只有在 num_layers > 1 时才生效。

例如：

nn.LSTM(
    input_size=100,
    hidden_size=128,
    num_layers=1,
    dropout=0.5,
    batch_first=True
)

这里 dropout=0.5 实际不会生效，因为只有 1 层。

如果是：

nn.LSTM(
    input_size=100,
    hidden_size=128,
    num_layers=2,
    dropout=0.5,
    batch_first=True
)

才会在 LSTM 层与层之间使用 dropout。

所以单层 LSTM 更推荐自己加：

self.dropout = nn.Dropout(0.5)

10. 双向 LSTM

双向 LSTM 写法：

self.lstm = nn.LSTM(
    input_size=embed_dim,
    hidden_size=hidden_size,
    num_layers=1,
    batch_first=True,
    bidirectional=True
)

双向后，输出维度会变成：

hidden_size * 2

因为它有：

正向 LSTM
反向 LSTM

所以分类层要改成：

self.classifier = nn.Linear(hidden_size * 2, num_classes)

BiLSTM 分类器

class BiLSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_classes, dropout=0.5):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embed_dim)

        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True,
            bidirectional=True
        )

        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden_size * 2, num_classes)

    def forward(self, x):
        embedded = self.embedding(x)

        output, (h_n, c_n) = self.lstm(embedded)

        # h_n shape: [2, batch_size, hidden_size]
        forward_hidden = h_n[-2]
        backward_hidden = h_n[-1]

        final_hidden = torch.cat(
            [forward_hidden, backward_hidden],
            dim=1
        )
        # final_hidden: [batch_size, hidden_size * 2]

        final_hidden = self.dropout(final_hidden)

        logits = self.classifier(final_hidden)

        return logits

11. BiLSTM 中 h_n 为什么这样取？

单层双向 LSTM 中：

h_n.shape = [2, batch_size, hidden_size]

其中：

h_n[0]：正向最后隐藏状态
h_n[1]：反向最后隐藏状态

也可以写成：

forward_hidden = h_n[0]
backward_hidden = h_n[1]

更通用一点写法是：

forward_hidden = h_n[-2]
backward_hidden = h_n[-1]

拼接：

final_hidden = torch.cat([forward_hidden, backward_hidden], dim=1)

最后得到：

[batch_size, hidden_size * 2]

12. LSTM 处理时间序列预测

如果不是文本，而是普通时间序列，比如：

过去 30 天的温度 → 预测明天温度

输入可以是：

x.shape = [batch_size, seq_len, feature_dim]

例如：

x.shape = [64, 30, 5]

含义：

64 个样本
每个样本 30 个时间步
每个时间步 5 个特征

模型：

class LSTMRegressor(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()

        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True
        )

        self.regressor = nn.Linear(hidden_size, 1)

    def forward(self, x):
        output, (h_n, c_n) = self.lstm(x)

        final_hidden = h_n[-1]

        pred = self.regressor(final_hidden)

        return pred

损失函数：

criterion = nn.MSELoss()

适合：

温度预测
销量预测
传感器预测
股票趋势预测

13. 常见维度错误

错误 1：输入少一维

错误输入：

x.shape = [batch_size, input_size]

LSTM 需要的是：

x.shape = [batch_size, seq_len, input_size]

如果只有一个时间步，也要写成：

x = x.unsqueeze(1)

变成：

[batch_size, 1, input_size]

错误 2：batch_first 没统一

如果你设置：

batch_first=True

那么输入必须是：

[batch_size, seq_len, input_size]

如果你设置：

batch_first=False

输入必须是：

[seq_len, batch_size, input_size]

这两个不要混。

错误 3：BiLSTM 后分类层维度没乘 2

如果使用：

bidirectional=True

那么：

final_hidden.shape = [batch_size, hidden_size * 2]

分类层必须是：

nn.Linear(hidden_size * 2, num_classes)

而不是：

nn.Linear(hidden_size, num_classes)

错误 4：CrossEntropyLoss 前手动 softmax

不要这样写：

probs = torch.softmax(logits, dim=1)
loss = criterion(probs, labels)

应该这样写：

loss = criterion(logits, labels)

因为 CrossEntropyLoss 内部已经包含了 log_softmax。

14. 第四讲核心总结

你要记住这几个形状：

x.shape = [batch_size, seq_len, input_size]
output.shape = [batch_size, seq_len, hidden_size]
h_n.shape = [num_layers, batch_size, hidden_size]
c_n.shape = [num_layers, batch_size, hidden_size]

做分类时常用：

final_hidden = h_n[-1]

然后：

logits = classifier(final_hidden)
# pytorch # lstm # 人工智能
Logo
AI Agent技术社区

Agent 垂直技术社区，欢迎活跃、内容共建。

更多推荐

cover

深入理解 AI Agent · AGENT #02：Agent 的四大核心机制，拆开来看每个都在干什么

avatar AI Agent技术社区

MCP协议开发实战：从零搭建AI Agent工具链

本文介绍了MCP协议在AI工具链中的应用与实践。MCP协议旨在解决异构工具集成难题，为AI模型提供标准化、安全的工具调用接口。文章详细解析了MCP的核心架构，包括Server、Client、Transport等组件，以及通信流程与安全模型。通过实战示例，指导读者开发计算器和文件系统工具Server，并集成到Claude等AI客户端中。进阶部分展示了如何构建复杂工具链，包括API集成、数据库查询和多

avatar AI Agent技术社区
cover

AI Tool 设计：粒度、参数与错误恢复怎么做

avatar AI Agent技术社区