mnist_savedmodel

Author: dskkato

examples/mnist_savedmodel.rs

@@ -37,10 +37,15 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // Load the saved model exported by regression_savedmodel.py.
     let mut graph = Graph::new();
-    let session =
-        SavedModelBundle::load(&SessionOptions::new(), &["serve"], &mut graph, export_dir)?.session;
-    let op_x = graph.operation_by_name_required("serving_default_sequential_input")?;
-    let op_predict = graph.operation_by_name_required("StatefulPartitionedCall")?;
+    let bundle =
+        SavedModelBundle::load(&SessionOptions::new(), &["serve"], &mut graph, export_dir)?;
+    let session = &bundle.session;
+
+    let signature = bundle.meta_graph_def().get_signature("serving_default")?;
+    let input_info = signature.get_input("input")?;
+    let op_x = graph.operation_by_name_required(&input_info.name().name)?;
+    let output_info = signature.get_output("output")?;
+    let op_predict = graph.operation_by_name_required(&output_info.name().name)?;
 
     // Train the model (e.g. for fine tuning).
     let mut args = SessionRunArgs::new();

examples/mnist_savedmodel/expected_values.txt

@@ -1 +1 @@
-4.0663176e-06, 1.4199884e-07, 9.556003e-05, 0.00065914105, 2.260991e-07, 4.076631e-06, 2.5459945e-09, 0.99904054, 1.5654963e-05, 0.00018059688
+3.112342e-05, 8.721303e-08, 0.0005018024, 0.0003709061, 1.6482764e-08, 1.8595395e-06, 1.3620006e-09, 0.999046, 5.4331244e-06, 4.2815118e-05

examples/mnist_savedmodel/mnist_savedmodel.py

@@ -31,7 +31,11 @@
 model.fit(x_train, y_train, epochs=1)
 
 # convert output type through softmax so that it can be interpreted as probability
-probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax(name="output")])
+inputs = tf.keras.Input((28, 28), name="input", dtype=tf.float32)
+x = model(inputs)
+outputs = tf.keras.layers.Softmax(name="output")(x)
+
+probability_model = tf.keras.Model(inputs=inputs, outputs=outputs)
 
 # dump expected values to compare Rust's outputs
 with open("examples/mnist_savedmodel/expected_values.txt", "w") as f: