use Keras's default signature

Author: dskkato

examples/mnist_savedmodel.rs

@@ -12,10 +12,6 @@ use tensorflow::Tensor;
 use image::io::Reader as ImageReader;
 use image::GenericImageView;
 
-#[cfg_attr(feature = "examples_system_alloc", global_allocator)]
-#[cfg(feature = "examples_system_alloc")]
-static ALLOCATOR: std::alloc::System = std::alloc::System;
-
 fn main() -> Result<(), Box<dyn Error>> {
     let export_dir = "examples/mnist_savedmodel"; // y = w * x + b
     if !Path::new(export_dir).exists() {
@@ -43,7 +39,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     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("predict_x")?;
+    let op_x = graph.operation_by_name_required("serving_default_sequential_input")?;
     let op_predict = graph.operation_by_name_required("StatefulPartitionedCall")?;
 
     // Train the model (e.g. for fine tuning).

examples/mnist_savedmodel/expected_values.txt

@@ -1 +1 @@
-2.9463042e-06, 1.12869834e-07, 0.00037566078, 0.00029264975, 2.2329447e-08, 1.9161714e-06, 1.8248756e-09, 0.999298, 2.3003201e-06, 2.6422202e-05
+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

examples/mnist_savedmodel/mnist_savedmodel.py

@@ -31,23 +31,20 @@
 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()])
-
-# convert keras model to TF2 function to get a computation graph
-x = tf.TensorSpec((None, 28, 28), tf.float32)
-predict = tf.function(lambda x: probability_model(x)).get_concrete_function(x=x)
+probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax(name="output")])
 
 # dump expected values to compare Rust's outputs
 with open("examples/mnist_savedmodel/expected_values.txt", "w") as f:
-    values = predict(x_test[:1, :, :]).numpy()
+    values = probability_model(x_test[:1, :, :])[0].numpy()
     print(*values, sep=", ", file=f)
 
 directory = "examples/mnist_savedmodel"
-signatures = {"predict": predict}
-tf.saved_model.save(model, directory, signatures=signatures)
+tf.saved_model.save(probability_model, directory)
 
 # export graph info to TensorBoard
 logdir = "logs/mnist_savedmodel"
 writer = tf.summary.create_file_writer(logdir)
+tf.summary.trace_on()
+values = probability_model(x_test[:1, :, :])
 with writer.as_default():
-    tf.summary.graph(predict.graph)
+    tf.summary.trace_export("Default", step=0)