如何将训练好的HDF5格式MNIST模型加载到Rust中实现预测

Rust 生态目前没有直接解析 Keras 原生 .h5 模型文件的成熟库，因此需要先在 Python 侧将模型转换为通用导出格式，再在 Rust 侧完成加载和预测操作，具体实现如下：

第一步：Python 侧转换模型格式 #

先把原有 .h5 模型转成适配 Rust 加载的格式，优先推荐转成轻量的 TFLite 格式，也可以转成 TensorFlow SavedModel 格式：

import tensorflow as tf

# 加载你训练好的h5模型
model = tf.keras.models.load_model("./sample_mnist.h5")

# 方案1：转成TFLite格式（推荐，体积小、加载快、依赖少）
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open("mnist_model.tflite", "wb") as f:
    f.write(tflite_model)

# 方案2：转成SavedModel格式（适配TensorFlow官方Rust绑定）
model.save("mnist_saved_model")

第二步：Rust 侧预测实现（TFLite方案，优先推荐） #

先配置 Cargo.toml 依赖 #

[dependencies]
tflite = "0.15.0"
ndarray = "0.15.6"

Rust 代码实现 #

use tflite::ops::builtin::BuiltinOpResolver;
use tflite::{FlatBufferModel, InterpreterBuilder};
use ndarray::Array4;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 加载转换后的TFLite模型
    let model = FlatBufferModel::build_from_file("mnist_model.tflite")?;
    let resolver = BuiltinOpResolver::default();
    let builder = InterpreterBuilder::new(model, resolver)?;
    let mut interpreter = builder.build()?;

    // 分配模型张量内存
    interpreter.allocate_tensors()?;

    // 构造输入：MNIST模型输入一般为[1, 28, 28, 1]的归一化浮点数组，根据你实际模型的输入格式调整
    // 此处替换为你实际的待预测输入数据
    let test_input: Array4<f32> = Array4::zeros((1, 28, 28, 1)); 

    // 把输入数据写入模型输入张量
    let input_tensor = interpreter.input_tensor_mut::<f32>(0)?;
    input_tensor.copy_from_slice(test_input.as_slice().unwrap());

    // 执行推理预测
    interpreter.invoke()?;

    // 获取输出结果，MNIST输出为10个数字对应的概率
    let output_tensor = interpreter.output_tensor::<f32>(0)?;
    // 取概率最高的索引作为预测结果
    let predicted_num = output_tensor
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
        .unwrap()
        .0;

    println!("预测的数字为：{}", predicted_num);
    Ok(())
}

备选方案：使用TensorFlow官方Rust绑定加载SavedModel #

如果需要使用完整TensorFlow能力，可以选择该方案：

Cargo.toml 依赖 #

[dependencies]
tensorflow = "0.20.0"

Rust 代码实现 #

use tensorflow::{Session, SessionOptions, Tensor, Graph, SavedModelBundle};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut graph = Graph::new();
    // 加载SavedModel
    let bundle = SavedModelBundle::load(
        &SessionOptions::new(),
        &["serve"],
        "mnist_saved_model",
        &mut graph,
    )?;

    // 构造输入张量，维度和数据类型和你训练时的输入保持一致
    let input = Tensor::<f32>::new(&[1, 28, 28, 1]); // 替换为实际输入数据

    // 运行推理，输入输出节点名可以在Python侧通过model.input_names、model.output_names查询后替换
    let mut outputs = bundle.session.run(
        &[("serving_default_input_1:0", &input)],
        &["StatefulPartitionedCall:0"],
        &[],
    )?;

    let output: Tensor<f32> = outputs.remove(0);
    let predicted_num = output
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
        .unwrap()
        .0;

    println!("预测的数字为：{}", predicted_num);
    Ok(())
}

注意事项 #

• 输入数据的维度、数据类型必须和你训练模型时的输入完全匹配，比如部分模型输入是扁平化的[1,784]格式，需要对应调整
• 不推荐直接解析.h5文件手动实现推理，该方案需要自己读取HDF5内的权重参数、手写每一层的计算逻辑，对复杂模型工作量极高

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