如何在Python中将TensorFlow模型存储为字符串而非冻结.pb文件

第一步：Python侧将模型序列化至内存字节流 #

不需要将冻结图写入本地磁盘，直接在内存中完成模型序列化，转码为可跨语言传递的字符串即可，示例代码如下：

import tensorflow as tf
import base64
from io import BytesIO

# 此处替换为你自己训练好的TensorFlow模型
model = tf.keras.applications.MobileNetV2(weights="imagenet", input_shape=(224,224,3))

# 定义模型推理签名
@tf.function(input_signature=[tf.TensorSpec([None, 224,224,3], tf.float32)])
def serving_fn(x):
    return {"predictions": model(x)}

# 将SavedModel直接导出到内存缓冲区，不落地本地文件
buf = BytesIO()
tf.saved_model.save(model, buf, signatures={"serving_default": serving_fn})

# 转base64编码字符串，适配跨语言接口传递规则
model_byte_str = base64.b64encode(buf.getvalue()).decode("utf-8")

# 直接调用你封装好的C++ DLL接口，传入上述字符串即可
# 示例DLL接口约定：load_model_from_str(const char* model_str)

第二步：C++侧DLL接收字符串并加载模型 #

接收Python传递的编码字符串后，先解码为原始二进制流，直接调用TensorFlow C++ API从内存加载模型，无需读取本地文件，示例代码如下：

#include <tensorflow/cc/saved_model/loader.h>
#include <tensorflow/cc/saved_model/tag_constants.h>
#include <string>
#include "base64.h" // 自行引入的base64解码工具头文件

tensorflow::SavedModelBundle model_bundle;
tensorflow::SessionOptions session_options;
tensorflow::RunOptions run_options;

// 导出给Python调用的模型加载接口
extern "C" __declspec(dllexport) int load_model_from_str(const char* model_base64_str) {
    // 1. 解码base64字符串为模型原始二进制数据
    std::string model_binary = base64_decode(model_base64_str);
    
    // 2. 直接从内存字节流加载模型
    tensorflow::SavedModelReader reader;
    auto status = reader.ReadFromBytes(tensorflow::Env::Default(), 
                                       {tensorflow::kSavedModelTagServe},
                                       model_binary, 
                                       &model_bundle);
    if (!status.ok()) {
        return -1; // 加载失败返回错误码
    }
    return 0;
}

// 后续训练、预测逻辑直接复用model_bundle.session运行计算即可，无需额外修改

注意事项 #

• 若模型体积较大，可直接传递原始二进制字节数组并约定长度参数，跳过base64编解码步骤，减少额外性能开销，同时避免字符串截断问题
• Python侧使用的TensorFlow版本必须和C++侧编译使用的TensorFlow版本完全一致，否则会出现模型兼容性问题
• 如果你使用的是旧版本单文件冻结图（.pb格式而非SavedModel格式），可以调整逻辑：Python侧直接序列化GraphDef为字节串，C++侧调用ParseFromArray接口直接从内存解析GraphDef，绑定到Session后即可正常运行

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