Keras：RNN仅输出最后值且存储全序列值用于模型解释的方法

阿华AIGC实验室

2026-5-27

嘿，这个需求我碰到过好多次，其实用框架的分支输出结构就能完美解决——既让后续层拿到最后一个时间步的输出，又能把全序列的RNN状态存下来做解释，而且完全不用重复计算RNN层，效率拉满。

用Keras（TensorFlow）的实现方式

Keras的函数式API天生适合这种多输出场景，步骤超简单：

先定义一个return_sequences=True的RNN层，这样它会输出整个时间序列的隐藏状态（形状是(batch_size, timesteps, units)）。
做两个分支：
- 一个分支用Lambda层提取最后一个时间步的输出，给后续的Dense层用；
- 另一个分支直接保留全序列的输出，专门用来做模型解释。
最后把输入和两个输出打包成一个多输出模型就行。

给你个具体的LSTM示例代码：

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, LSTM, Dense, Lambda

# 假设你的输入是 (时间步长, 特征数)
timesteps = 10
features = 5

# 定义输入层
input_layer = Input(shape=(timesteps, features))

# 共享的LSTM层，输出全序列
lstm_full = LSTM(64, return_sequences=True)(input_layer)

# 分支1：提取最后一个时间步的输出，用于后续Dense层
last_step_out = Lambda(lambda x: x[:, -1, :])(lstm_full)
final_pred = Dense(1, activation='sigmoid')(last_step_out)

# 分支2：保留全序列的LSTM输出，用于模型解释
sequence_log = lstm_full

# 构建多输出模型
model = Model(inputs=input_layer, outputs=[final_pred, sequence_log])

# 编译的时候，全序列输出不需要计算损失，设为None就行
model.compile(optimizer='adam', loss=['binary_crossentropy', None])

训练的时候，你只需要给final_pred对应的标签就行；推理的时候，调用model.predict(x)会返回两个结果：第一个是最终的预测值，第二个就是整个序列的LSTM隐藏状态，直接拿去做解释就好。

用PyTorch的实现方式

PyTorch的话，直接在模型的forward函数里返回两个值就行，逻辑更直观：

import torch
import torch.nn as nn

class RNNInterpretable(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        # output: (batch_size, timesteps, hidden_size) —— 全序列状态
        # h_n: (1, batch_size, hidden_size) —— 最后一个时间步的隐藏状态
        output, (h_n, _) = self.lstm(x)
        # 把h_n的维度压缩成 (batch_size, hidden_size)
        last_hidden = h_n.squeeze(0)
        # 计算最终预测
        pred = self.fc(last_hidden)
        # 同时返回预测值和全序列状态
        return pred, output

使用的时候，训练阶段只需要用到pred来计算损失，output直接存下来或者忽略；推理的时候，拿到两个返回值，output就是你需要的全序列解释数据。