tf.train.write_graph被Java调用时无法保存模型，移动端模型保存求助

咱们先拆解你遇到的问题，从代码本身到Java/Android调用的权限，一步步梳理清楚：

1. 代码本身的核心问题 #

你的代码存在几个明显的逻辑错误，直接导致模型无法保存：

• 保存函数未执行：你在tf.Session()上下文里定义了save_model()函数，但从头到尾没调用它，自然不会触发保存操作。
• 变量未初始化：W和b是tf.Variable类型，但你没有在Session中执行变量初始化（比如tf.global_variables_initializer().run()），直接保存的graph_def会缺失变量的实际值，甚至可能导致图结构不完整。
• 无意义的错误代码：return tf.to_float(tf.rank(path))完全不合理——path是字符串路径，tf.rank是用来计算张量维度的API，对字符串调用会直接抛出错误，这部分代码可以直接删掉。
• 图构建顺序混乱：你在Session上下文外定义占位符和变量，又在Session里写保存函数，这种写法容易导致图节点绑定错误，应该先完成所有图结构的构建，再启动Session执行后续操作。

2. Java/Android调用时的权限问题 #

Android应用的/data/data/com.chelexa.tfandroid是私有目录，但通过JNI调用Python代码时，需要注意：

• 确保Python进程拥有该目录的写入权限（私有目录默认是允许的，但如果路径是外部存储，需要在AndroidManifest.xml中申请WRITE_EXTERNAL_STORAGE权限）。
• 提前确认目录存在，可以在Python代码中用os.makedirs自动创建目录，避免因目录不存在导致保存失败。

修正后的可运行代码 #

下面是调整后的代码，不仅能正确保存模型，还生成了适合移动端使用的冻结图（将变量值直接嵌入图中，无需额外加载checkpoint）：

import tensorflow as tf
import os

# 第一步：先完成所有图结构的构建，不要在Session内定义节点或函数
x = tf.placeholder(tf.float32, shape=[None, 784], name="x")
y = tf.placeholder(tf.float32, [None, 10], name="y")
W = tf.Variable(tf.zeros([784, 10]), name="weights")
b = tf.Variable(tf.zeros([10]), name="biases")
# 必须给输出节点命名，方便移动端调用
logits = tf.matmul(x, W) + b
predictions = tf.nn.softmax(logits, name="predictions")

# 定义保存模型的工具函数
def save_model(sess, save_dir, model_name):
    # 确保保存目录存在，不存在则创建
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    
    # 保存原始graph_def（仅图结构，不含变量值）
    graph_path = tf.train.write_graph(sess.graph_def, save_dir, f"{model_name}.pb", as_text=False)
    print(f"图结构已保存至: {graph_path}")
    
    # 生成适合移动端的冻结图（将变量转为常量，嵌入图中）
    from tensorflow.python.framework import graph_util
    frozen_graph_def = graph_util.convert_variables_to_constants(
        sess,
        sess.graph_def,
        output_node_names=["predictions"]  # 指定移动端需要调用的输出节点名
    )
    # 写入冻结图文件
    frozen_path = os.path.join(save_dir, f"{model_name}_frozen.pb")
    with open(frozen_path, "wb") as f:
        f.write(frozen_graph_def.SerializeToString())
    print(f"冻结图已保存至: {frozen_path}")
    return frozen_path

# 第二步：启动Session，初始化变量，执行保存
with tf.Session() as sess:
    # 初始化所有变量
    sess.run(tf.global_variables_initializer())
    
    # 这里可以添加你的模型训练代码（如果需要先训练再保存）
    # ...
    
    # 调用保存函数
    save_model(sess, "/data/data/com.chelexa.tfandroid", "mnist_50_mlp")

额外注意事项 #

• 优先使用冻结图：tf.train.write_graph只保存图结构，移动端加载后还需要额外加载变量checkpoint，非常繁琐。上面代码生成的冻结图直接包含变量值，移动端可以直接加载使用。
• Java调用时的路径验证：在JNI调用Python脚本前，可以先在Python代码中打印os.path.abspath(save_dir)，确认路径是否正确、是否存在。
• 避免Session上下文混乱：所有图节点的定义都要放在Session启动前，Session只用来执行初始化、训练、保存等操作。

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