如何在TensorFlow中从计算图将RNN状态保存到文件？

刚好之前处理过类似的TensorFlow RNN状态保存+Serving的场景，结合你给出的TensorFlow Serving API代码片段，我来一步步拆解怎么做：

一、先在模型构建阶段搞定状态的可追踪性 #

RNN的状态是跨批次的“记忆”，要能保存它，首先得让TensorFlow能追踪到这个状态——要么用tf.Variable显式定义，要么用Keras的有状态RNN：

• 原生TensorFlow写法：把RNN初始状态定义为可训练/可保存的变量，这样它会自动被纳入SavedModel的变量集合：

# 定义RNN状态变量，batch_size和state_size根据你的模型调整
rnn_state = tf.Variable(tf.zeros([batch_size, state_size]), name="rnn_state")
# 构建RNN计算图
output, new_state = tf.nn.dynamic_rnn(rnn_cell, inputs, initial_state=rnn_state)
# 生成更新状态的操作
update_state_op = tf.assign(rnn_state, new_state)

• Keras有状态RNN写法：给RNN/LSTM层加stateful=True，Keras会自动帮你管理状态变量，导出SavedModel时也会自动包含这些状态：

lstm_layer = tf.keras.layers.LSTM(64, stateful=True, return_state=True)
output, h_state, c_state = lstm_layer(inputs)

二、导出SavedModel时要包含状态的访问逻辑 #

导出SavedModel的时候，得把状态的输出（或者更新逻辑）加入签名，这样Serving端才能拿到最新的状态：

export_dir = "./my_rnn_saved_model"

# 定义带状态输出的预测函数
@tf.function(input_signature=[tf.TensorSpec(shape=[None, seq_len, feature_dim], dtype=tf.float32)])
def predict_with_state(inputs):
    # 假设model是你构建好的RNN模型，return_state=True会返回更新后的状态
    output, new_h, new_c = model(inputs, training=False)
    # 返回预测输出和最新状态
    return {"pred_output": output, "new_h_state": new_h, "new_c_state": new_c}

# 导出SavedModel，把这个函数作为默认签名
tf.saved_model.save(model, export_dir, signatures={"serving_default": predict_with_state})

三、在TensorFlow Serving的Predict逻辑中保存状态到文件 #

结合你给出的SavedModelPredict代码片段，我们可以在处理完预测请求后，把最新的RNN状态序列化写入文件：

Status SavedModelPredict(const RunOptions& run_options, ServerCore* core,
                         const PredictRequest& request, PredictResponse* response) {
  // 获取模型的Servable Handle
  ServableHandle<SavedModelBundle> bundle;
  TF_RETURN_IF_ERROR(core->GetServableHandle(request.model_spec(), &bundle));
  
  // 确定要使用的签名名称，默认用"serving_default"
  const string signature_name = request.model_spec().signature_name().empty() 
      ? kDefaultServingSignatureDefName 
      : request.model_spec().signature_name();
  const SignatureDef& signature = bundle->meta_graph_def.signature_def().at(signature_name);
  
  auto session = bundle->session.get();
  std::vector<std::pair<string, Tensor>> inputs;
  // 这里需要解析PredictRequest里的输入数据，填充到inputs中，比如：
  // inputs.emplace_back(signature.inputs().at("inputs").name(), input_tensor);
  
  // 定义要获取的输出：预测结果+最新状态
  std::vector<string> output_names = {"pred_output", "new_h_state", "new_c_state"};
  std::vector<Tensor> outputs;
  
  // 运行计算图，拿到结果和状态
  TF_RETURN_IF_ERROR(session->Run(run_options, inputs, output_names, {}, &outputs));
  
  // 把状态张量序列化后写入文件
  // 以隐藏层状态为例，你也可以同时保存细胞状态
  const Tensor& h_state_tensor = outputs[1];
  TensorProto state_proto;
  h_state_tensor.AsProtoTensorContent(&state_proto);
  
  std::ofstream state_file("./rnn_h_state.pb", std::ios::binary);
  if (!state_file.write(state_proto.SerializeAsString().data(), state_proto.ByteSizeLong())) {
    return Status(error::Code::INTERNAL, "Failed to write RNN hidden state to file");
  }
  
  // 填充PredictResponse的预测结果部分
  // 比如把outputs[0]的内容写入response->mutable_outputs()->at("pred_output")...
  
  return Status::OK();
}

几个要注意的点 #

• 状态形状一致性：保存和恢复状态时，要确保张量的形状完全匹配，比如batch_size如果是动态的，可能需要用tf.Variable.assign来兼容不同批次的形状。
• 并发安全：如果你的Serving服务处理多并发请求，直接写文件会有冲突，建议要么为每个请求维护独立的状态（无状态服务），要么加文件锁或者用分布式存储来管理状态。
• 状态恢复：如果之后要恢复状态，可以反向操作：读取文件中的TensorProto，解析成Tensor，然后通过Session的Run操作赋值给模型的状态变量。

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