Improve SavedModelBundle support

- Give explicit name to an operation in regression_savedmodel script
- Add test_resources with a saved model
- Add test for saved model bundle
- Fix typo
- Wipe try! macros

Author: Enet4

examples/regression_savedmodel.py

@@ -10,7 +10,7 @@
 
 w = tf.Variable(tf.random_uniform([1], -1.0, 1.0), name='w')
 b = tf.Variable(tf.zeros([1]), name='b')
-y_hat = w * x + b
+y_hat = tf.add(w * x, b, name="y_hat")
 
 loss = tf.reduce_mean(tf.square(y_hat - y))
 optimizer = tf.train.GradientDescentOptimizer(0.5)

src/session.rs

@@ -22,7 +22,7 @@ use super::TensorType;
 pub struct SavedModelBundle {
     /// The loaded session.
     pub session: Session,
-    /// A meta graph defition as raw protocol buffer.
+    /// A meta graph definition as raw protocol buffer.
     pub meta_graph_def: Vec<u8>,
 }
 
@@ -38,15 +38,15 @@ impl SavedModelBundle {
         let mut status = Status::new();
 
         let export_dir_cstr =
-            try!(export_dir.as_ref()
+            export_dir.as_ref()
                      .to_str()
                      .and_then(|s| CString::new(s.as_bytes()).ok())
-                     .ok_or_else(|| invalid_arg!("Invalid export directory path")));
+                     .ok_or_else(|| invalid_arg!("Invalid export directory path"))?;
 
-        let tags_cstr: Vec<_> = try!(tags.into_iter()
-                                         .map(|t| CString::new(t.as_ref()))
-                                         .collect::<::std::result::Result<_, _>>()
-                                         .map_err(|_| invalid_arg!("Invalid tag name")));
+        let tags_cstr: Vec<_> = tags.into_iter()
+                                    .map(|t| CString::new(t.as_ref()))
+                                    .collect::<::std::result::Result<_, _>>()
+                                    .map_err(|_| invalid_arg!("Invalid tag name"))?;
         let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();
 
         // The empty TF_Buffer will be filled by LoadSessionFromSavedModel
@@ -102,16 +102,15 @@ impl Session {
          -> Result<Self> {
         let mut status = Status::new();
 
-        let export_dir_cstr =
-            try!(export_dir.as_ref()
+        let export_dir_cstr = export_dir.as_ref()
                      .to_str()
                      .and_then(|s| CString::new(s.as_bytes()).ok())
-                     .ok_or_else(|| invalid_arg!("Invalid export directory path")));
+                     .ok_or_else(|| invalid_arg!("Invalid export directory path"))?;
 
-        let tags_cstr: Vec<_> = try!(tags.into_iter()
-                                         .map(|t| CString::new(t.as_ref()))
-                                         .collect::<::std::result::Result<_, _>>()
-                                         .map_err(|_| invalid_arg!("Invalid tag name")));
+        let tags_cstr: Vec<_> = tags.into_iter()
+                                    .map(|t| CString::new(t.as_ref()))
+                                    .collect::<::std::result::Result<_, _>>()
+                                    .map_err(|_| invalid_arg!("Invalid tag name"))?;
         // keeping tags_cstr to retain strings in memory
         let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();
 
@@ -385,4 +384,39 @@ mod tests {
         assert_eq!(output_tensor[0], 4.0);
         assert_eq!(output_tensor[1], 6.0);
     }
+
+    #[test]
+    fn test_savedmodelbundle() {
+        let mut graph = Graph::new();
+        let bundle = SavedModelBundle::load(
+            &SessionOptions::new(),
+            &["train", "serve"],
+            &mut graph,
+            "test_resources/regression-model",
+        ).unwrap();
+
+        let x_op = graph.operation_by_name_required("x").unwrap();
+        let y_op = graph.operation_by_name_required("y").unwrap();
+        let y_hat_op = graph.operation_by_name_required("y_hat").unwrap();
+        let _train_op = graph.operation_by_name_required("train").unwrap();
+
+        let SavedModelBundle {
+            mut session,
+            meta_graph_def,
+        } = bundle;
+
+        assert!(!meta_graph_def.is_empty());
+
+        let mut x = <Tensor<f32>>::new(&[1]);
+        x[0] = 2.0;
+        let mut y = <Tensor<f32>>::new(&[1]);
+        y[0] = 4.0;
+        let mut step = StepWithGraph::new();
+        step.add_input(&x_op, 0, &x);
+        step.add_input(&y_op, 0, &y);
+        let output_token = step.request_output(&y_hat_op, 0);
+        session.run(&mut step).unwrap();
+        let output_tensor = step.take_output::<f32>(output_token).unwrap();
+        assert_eq!(output_tensor.len(), 1);
+    }
 }