将经过训练的HDF5模型加载到Rust中以进行预测

Question

我使用MNIST数据集训练了一个模型来识别数字。该模型已经使用TensorFlow和Keras在Python语言中进行了训练，并将输出保存到我命名为"sample_mnist.h5“的HDF5文件中。

我想将训练好的模型从HDF5文件加载到Rust中进行预测。

在Python语言中，我可以从HDF5生成模型，并使用代码进行预测：

model = keras.models.load_model("./sample_mnist.h5")
model.precict(test_input)  # assumes test_input is the correct input type for the model

这段Python代码片段的Rust等价物是什么？

Answer 1

首先，您需要将模型保存为.pb格式，而不是.hdf5格式，以便将其移植到Rust，因为此格式保存了所有关于在Python外重建它所需的模型执行图的。在TensorFlow Rust存储库上有一个来自用户justnoxx的打开的pull request，它展示了如何为一个简单的模型执行此操作。要点是给定Python中的一些模型...

from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense

classifier = Sequential()
classifier.add(Dense(5, activation='relu', name="test_in", input_dim=5)) # Named input
classifier.add(Dense(5, activation='relu'))
classifier.add(Dense(1, activation='sigmoid', name="test_out")) # Named output

classifier.compile(optimizer ='adam', loss='binary_crossentropy', metrics=['accuracy'])

classifier.fit([[0.1, 0.2, 0.3, 0.4, 0.5]], [[1]], batch_size=1, epochs=1);

classifier.save('examples/keras_single_input_saved_model', save_format='tf')

和我们命名的输入"test_in“和输出"test_out”以及它们的预期大小，我们可以将保存的模型应用于Rust ...

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

fn main() {

    // In this file test_in_input is being used while in the python script,
    // that generates the saved model from Keras model it has a name "test_in".
    // For multiple inputs _input is not being appended to signature input parameter name.
    let signature_input_parameter_name = "test_in_input";
    let signature_output_parameter_name = "test_out";

    // Initialize save_dir, input tensor, and an empty graph
    let save_dir =
        "examples/keras_single_input_saved_model";
    let tensor: Tensor<f32> = Tensor::new(&[1, 5])
        .with_values(&[0.1, 0.2, 0.3, 0.4, 0.5])
        .expect("Can't create tensor");
    let mut graph = Graph::new();

    // Load saved model bundle (session state + meta_graph data)
    let bundle = 
        SavedModelBundle::load(&SessionOptions::new(), &["serve"], &mut graph, save_dir)
        .expect("Can't load saved model");

    // Get the session from the loaded model bundle
    let session = &bundle.session;

    // Get signature metadata from the model bundle
    let signature = bundle
        .meta_graph_def()
        .get_signature("serving_default")
        .unwrap();

    // Get input/output info
    let input_info = signature.get_input(signature_input_parameter_name).unwrap();
    let output_info = signature
        .get_output(signature_output_parameter_name)
        .unwrap();

    // Get input/output ops from graph
    let input_op = graph
        .operation_by_name_required(&input_info.name().name)
        .unwrap();
    let output_op = graph
        .operation_by_name_required(&output_info.name().name)
        .unwrap();
    
    // Manages inputs and outputs for the execution of the graph
    let mut args = SessionRunArgs::new();
    args.add_feed(&input_op, 0, &tensor); // Add any inputs

    let out = args.request_fetch(&output_op, 0); // Request outputs

    // Run model
    session.run(&mut args) // Pass to session to run
        .expect("Error occurred during calculations");

    // Fetch outputs after graph execution
    let out_res: f32 = args.fetch(out).unwrap()[0];

    println!("Results: {:?}", out_res);
}