如何在Keras或PyTorch中获取LSTM各时间步细胞态用于循环网络可视化

Keras 实现（TF 2.x / 原生Keras通用） #

核心原理 #

无需重写LSTM计算逻辑，直接复用已训练LSTM层内置的LSTMCell实例，通过官方提供的backend.rnn循环在每一步收集细胞态c，全程不改动框架源码，完全复用已训练的权重。

操作步骤 #

1. 原有训练流程无需修改，仅需给Embedding、LSTM层添加明确命名，方便后续获取层实例：

   inputs = Input(shape=(MAX_LEN,))
   embd = Embedding(input_dim=MAX_FEATURES, output_dim=128, input_length=MAX_LEN, name='embd')(inputs)
   lstm = LSTM(units=128, return_sequences=True, name='lstm')(embd)
   flt = Flatten()(lstm)
   drp = Dropout(0.5)(flt)
   dns = Dense(N_CLASSES)(drp)
   activ = Activation('softmax')(dns)
   training_model = Model([inputs], [activ])
   training_model.compile(loss='categorical_crossentropy', optimizer='rmsprop')
   training_model.fit(X_train, y_train)
   

2. 训练完成后，基于已训练模型提取逐时间步的细胞态：

   import tensorflow as tf # 原生Keras可替换为import keras.backend as K
   # 取出已训练好的LSTM层和绑定的内置LSTMCell
   trained_lstm = training_model.get_layer('lstm')
   lstm_cell = trained_lstm.cell
   # 先获取Embedding层的输出（即原LSTM层的输入）
   embd_infer_model = Model(inputs=training_model.input, outputs=training_model.get_layer('embd').output)
   embd_output = embd_infer_model.predict(X_test) # shape: [batch, MAX_LEN, 128]
   # 定义单步计算逻辑，直接调用内置LSTMCell，仅收集每步的c态
   def step(carry, current_input):
       h_prev, c_prev = carry
       # 复用已训练的门控计算逻辑，无自定义实现
       h_new, [c_new] = lstm_cell(current_input, states=[h_prev, c_prev])
       return (h_new, c_new), c_new
   # 初始化和原LSTM默认逻辑一致的全零初始状态
   batch_size = tf.shape(embd_output)[0]
   initial_state = [
       tf.zeros((batch_size, 128)),
       tf.zeros((batch_size, 128))
   ]
   # 调用官方rnn循环，收集所有时间步的c态
   cell_state_seq, _ = tf.keras.backend.rnn(
       step,
       inputs=embd_output,
       initial_states=initial_state,
       go_backwards=False,
       mask=None,
       unroll=False
   )
   # 最终cell_state_seq shape: [batch, MAX_LEN, 128]，即每个时间步的细胞态
   # 直接对输出做*tanh变换*即可得到论文图2要求的可视化输入
   c_visual = tf.tanh(cell_state_seq)
   

注意事项 #

• 若训练时使用了自定义的初始隐状态/细胞态，需同步替换代码中的initial_state取值，保证和训练逻辑一致
• 原生Keras用户仅需将tf.keras.backend.rnn替换为keras.backend.rnn，其余逻辑完全通用
• 无需重新训练模型，所有操作均在推理阶段完成

PyTorch 实现（适配新手） #

核心原理 #

无需重写LSTM门控计算逻辑，将已训练nn.LSTM的权重直接拷贝到官方提供的nn.LSTMCell实例中，按时间步循环计算时直接收集每步的c态，全程不改动框架源码，完全复用已训练权重。

操作步骤 #

1. 假设你已经完成了LSTM模型的训练，且模型结构如下：

   import torch
   import torch.nn as nn
   device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
   class TrainedLSTM(nn.Module):
       def __init__(self, max_features, embed_dim, hidden_dim, n_classes, max_len):
           super().__init__()
           self.embedding = nn.Embedding(max_features, embed_dim)
           self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
           self.fc = nn.Linear(hidden_dim * max_len, n_classes)
           self.softmax = nn.Softmax(dim=-1)
       def forward(self, x):
           x = self.embedding(x)
           output, (h_n, c_n) = self.lstm(x)
           x = output.flatten(1)
           return self.softmax(self.fc(x))
   # 模型已完成训练，权重加载完毕
   trained_model = TrainedLSTM(MAX_FEATURES, 128, 128, N_CLASSES, MAX_LEN).to(device)
   

2. 基于已训练模型提取逐时间步的细胞态：

   # 第一步：将已训练LSTM的权重拷贝到独立的LSTMCell实例中（权重布局完全兼容）
   lstm_cell = nn.LSTMCell(input_size=128, hidden_size=128).to(device)
   lstm_cell.weight_ih.data = trained_model.lstm.weight_ih_l0.data
   lstm_cell.weight_hh.data = trained_model.lstm.weight_hh_l0.data
   lstm_cell.bias_ih.data = trained_model.lstm.bias_ih_l0.data
   lstm_cell.bias_hh.data = trained_model.lstm.bias_hh_l0.data
   # 第二步：按时间步循环收集c态
   def get_all_cell_states(x_test):
       trained_model.eval()
       all_c_states = []
       with torch.no_grad(): # 推理阶段关闭梯度计算，节省内存
           embed_out = trained_model.embedding(x_test.to(device)) # shape: [batch, MAX_LEN, 128]
           batch_size, seq_len, _ = embed_out.shape
           # 初始化和原LSTM默认逻辑一致的全零初始状态
           h_t = torch.zeros(batch_size, 128).to(device)
           c_t = torch.zeros(batch_size, 128).to(device)
           # 遍历每个时间步计算
           for t in range(seq_len):
               x_t = embed_out[:, t, :]
               # 调用内置LSTMCell计算，复用已拷贝的训练权重
               h_t, c_t = lstm_cell(x_t, (h_t, c_t))
               all_c_states.append(c_t.unsqueeze(1))
       # 拼接所有时间步的c态，shape: [batch, MAX_LEN, 128]
       return torch.cat(all_c_states, dim=1)
   # 调用获取c态，直接对输出做*tanh变换*得到论文要求的可视化输入
   all_c = get_all_cell_states(X_test)
   c_visual = torch.tanh(all_c)
   

注意事项 #

• 若使用多层LSTM，仅需对应每层将权重拷贝到独立的LSTMCell实例，按层顺序遍历计算即可，逻辑完全一致
• 全程无自定义LSTM计算逻辑，门控运算完全由官方API实现，计算结果和原nn.LSTM内部逻辑完全一致
• 无需重新训练模型，所有操作均在推理阶段完成

内容的提问来源于stack exchange，提问作者NFa