TensorFlow 2.6协议缓冲区加载Java 8遇解析错误求解决方案

报错原因解析 #

出现While parsing a protocol message, the input ended unexpectedly in the middle of the field.通常有几个常见诱因：

• 模型文件损坏：导出过程中断、文件传输时截断都会导致.pb文件不完整，先核对Python端导出的文件大小和Java端接收的是否一致，重新导出一次试试
• 导出方式不当：如果用Keras的model.save()导出，模型会包含Keras专属元数据，Java端解析时容易出错；非Keras模型必须用tf.saved_model.save()并明确指定输入输出签名
• 输入签名不匹配：Java端加载时传入的张量形状、类型和模型导出时定义的签名不一致，也可能触发解析错误

非Keras TF 2.x模型导出与Java加载示例 #

Python端导出Transformer模型 #

假设你的Transformer是用原生TensorFlow API构建的（未继承Keras Model），按以下步骤导出：

import tensorflow as tf
# 导入你自己实现的Transformer类
from your_transformer_code import Transformer

# 初始化Transformer实例
transformer = Transformer(
    num_layers=6, d_model=512, num_heads=8, dff=2048,
    input_vocab_size=8500, target_vocab_size=8000
)

# 定义输入签名，要和实际推理时的输入形状、类型完全一致
input_specs = [
    tf.TensorSpec(shape=(None, None), dtype=tf.int32, name="encoder_input"),
    tf.TensorSpec(shape=(None, None), dtype=tf.int32, name="decoder_input")
]

# 用tf.function包装推理逻辑，绑定输入签名
@tf.function(input_signature=input_specs)
def run_inference(encoder_input, decoder_input):
    output, _ = transformer(
        encoder_input, decoder_input,
        training=False,
        encoder_mask=None,
        decoder_mask=None,
        look_ahead_mask=None
    )
    return {"output": output}

# 保存为SavedModel格式（包含.pb和变量文件的目录）
tf.saved_model.save(transformer, "./transformer_model")

注意：不要单独提取.pb文件，Java端需要加载整个SavedModel目录，单独的.pb缺少变量等依赖文件。

Java端加载并推理 #

使用TensorFlow Java API（Maven依赖请用对应版本，比如org.tensorflow:tensorflow:2.6.0，尽量和Python端TF版本一致）：

import org.tensorflow.SavedModelBundle;
import org.tensorflow.Tensor;
import org.tensorflow.types.TInt32;
import org.tensorflow.ndarray.Shape;
import org.tensorflow.ndarray.IntNdArray;
import org.tensorflow.ndarray.NdArrays;

public class TransformerJavaInference {
    public static void main(String[] args) {
        // 加载SavedModel目录，"serve"是默认的签名标签
        try (SavedModelBundle model = SavedModelBundle.load("./transformer_model", "serve")) {
            // 构造示例输入：encoder输入为[[1,2,3]], decoder输入为[[4,5,6]]
            IntNdArray encoderInputArr = NdArrays.ofInts(Shape.of(1, 3));
            encoderInputArr.set(1, 0, 0).set(2, 0, 1).set(3, 0, 2);
            Tensor<TInt32> encoderTensor = TInt32.tensorOf(encoderInputArr);

            IntNdArray decoderInputArr = NdArrays.ofInts(Shape.of(1, 3));
            decoderInputArr.set(4, 0, 0).set(5, 0, 1).set(6, 0, 2);
            Tensor<TInt32> decoderTensor = TInt32.tensorOf(decoderInputArr);

            // 运行推理，输入名称要和Python端定义的完全匹配
            Tensor<?> outputTensor = model.session().runner()
                    .feed("encoder_input", encoderTensor)
                    .feed("decoder_input", decoderTensor)
                    .fetch("output")
                    .run()
                    .get(0);

            // 解析输出结果
            float[][] output = outputTensor.copyTo(new float[1][3]);
            System.out.println("推理输出：");
            for (float[] row : output) {
                for (float val : row) {
                    System.out.print(val + " ");
                }
                System.out.println();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

额外注意点 #

• 版本兼容性：TensorFlow Python和Java版本尽量保持同大版本，避免因版本差异导致的解析错误
• 输入签名必须明确：导出时的input_signature要准确对应推理时的输入，Java端传入的张量形状、类型要和签名一致
• 必须加载完整SavedModel目录：单独的.pb文件无法包含模型的所有变量和元数据，Java端加载会失败

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