如何在Java中加载TensorFlow的DNNClassifier模型及Python加载报错咨询

我来帮你梳理下在Java中加载这个TensorFlow SavedModel的完整流程，亲测可行，步骤如下：

1. 引入TensorFlow Java依赖 #

首先得给你的Java项目加上TensorFlow的依赖，版本尽量和你训练模型时用的Python TensorFlow版本保持一致（比如都是2.x系列），避免兼容性问题。

• 如果用Maven，在pom.xml中添加：

<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>tensorflow</artifactId>
    <version>2.15.0</version>
</dependency>

• 如果用Gradle，在build.gradle中添加：

implementation 'org.tensorflow:tensorflow:2.15.0'

2. 确认模型的输入输出签名 #

这一步非常关键！DNNClassifier导出的SavedModel有固定的签名规则，你得先搞清楚输入输出张量的具体名称。可以在Python环境下用命令行工具查看：

saved_model_cli show --dir exported_path --all

输出里的signature_def板块会显示详细信息，比如默认预测签名是serving_default，输入可能叫input或者dnn/input_from_feature_columns/input_layer，输出通常包含class_ids（预测的类别ID）、probabilities（每个类别的概率）等。把这些名称记下来，后面代码要用到。

3. 完整的Java预测代码示例 #

假设我们从签名里得到：

• 输入张量名称："input"（实际以你的签名输出为准）
• 输出张量名称："class_ids"和"probabilities"

代码如下：

import org.tensorflow.SavedModelBundle;
import org.tensorflow.Tensor;
import org.tensorflow.Tensors;
import java.util.Arrays;

public class DNNClassifierPredictor {
    public static void main(String[] args) {
        // 替换成你的模型实际路径
        String modelDir = "你的模型路径/exported_path";
        
        // 加载模型，"serve"是默认的服务签名标签
        try (SavedModelBundle model = SavedModelBundle.load(modelDir, "serve")) {
            // 准备输入数据：32个浮点数的特征数组
            float[] inputFeatures = new float[32];
            // 这里替换成你的实际输入数据，比如inputFeatures[0] = 1.2f; 按需赋值
            Arrays.fill(inputFeatures, 0.6f); // 示例填充数据
            
            // 构造输入Tensor：形状是[1, 32]，代表1个样本，每个样本32个特征
            try (Tensor<Float> inputTensor = Tensors.create(new long[]{1, 32}, inputFeatures)) {
                // 运行预测，指定输入和要获取的输出
                var outputs = model.session().runner()
                        .feed("input", inputTensor) // 输入名称要和签名一致
                        .fetch("class_ids") // 获取预测的类别ID
                        .fetch("probabilities") // 获取类别概率分布
                        .run();
                
                // 解析预测的类别ID
                try (Tensor<Long> classIdTensor = outputs.get(0).expect(Long.class)) {
                    long[] classIds = new long[1];
                    classIdTensor.copyTo(classIds);
                    System.out.println("预测类别ID：" + classIds[0]);
                }
                
                // 解析类别概率分布
                try (Tensor<Float> probTensor = outputs.get(1).expect(Float.class)) {
                    float[] probabilities = new float[4];
                    probTensor.copyTo(probabilities);
                    System.out.println("4个类别概率：" + Arrays.toString(probabilities));
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

4. 关键注意事项 #

• 张量形状匹配：输入张量的形状必须和训练时一致，这里是[1,32]（批量大小1，特征数32），如果要批量预测，可改成[N,32]（N为样本数量）。
• 资源自动释放：Tensor、SavedModelBundle这些资源要放在try-with-resources块中，确保使用后自动释放，避免内存泄漏。
• 签名名称核对：如果运行时提示找不到张量，一定是输入输出名称和签名不匹配，再用saved_model_cli命令仔细核对。

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