move from_saved_model_to_bundle to SavedModelBundle::load

Author: Enet4

src/session.rs

@@ -26,31 +26,15 @@ pub struct SavedModelBundle {
     pub meta_graph_def: Vec<u8>,
 }
 
-/// Manages a single graph and execution.
-#[derive(Debug)]
-pub struct Session {
-    inner: *mut tf::TF_Session,
-}
-
-impl Session {
-    /// Creates a session.
-    pub fn new(options: &SessionOptions, graph: &Graph) -> Result<Self> {
-        let mut status = Status::new();
-        let inner = unsafe { tf::TF_NewSession(graph.inner(), options.inner, status.inner()) };
-        if inner.is_null() {
-            Err(status)
-        } else {
-            Ok(Session { inner: inner })
-        }
-    }
+impl SavedModelBundle {
 
-    /// Loads a session from an exported model.
-    pub fn from_saved_model<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
+    /// Loads a session from an exported model, creating a bundle
+    pub fn load<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
         (options: &SessionOptions,
          tags: Tags,
          graph: &mut Graph,
          export_dir: P)
-         -> Result<Self> {
+         -> Result<SavedModelBundle> {
         let mut status = Status::new();
 
         let export_dir_cstr =
@@ -63,33 +47,59 @@ impl Session {
                                          .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();
 
+        // The empty TF_Buffer will be filled by LoadSessionFromSavedModel
+        let mut meta = unsafe { Buffer::<u8>::from_ptr(ptr::null_mut(), 0) };
+
         let inner = unsafe {
             tf::TF_LoadSessionFromSavedModel(options.inner,
                                              ptr::null(),
                                              export_dir_cstr.as_ptr(),
                                              tags_ptr.as_ptr(),
                                              tags_ptr.len() as c_int,
                                              graph.inner(),
-                                             ptr::null_mut(),
+                                             meta.inner_mut(),
                                              status.inner())
         };
+        if inner.is_null() {
+            Err(status)
+        } else {
+            let session = Session { inner: inner };
+            Ok(SavedModelBundle {
+                session: session,
+                meta_graph_def: Vec::from(meta.as_ref())
+            })
+        }
+    }
+    
+}
+
+/// Manages a single graph and execution.
+#[derive(Debug)]
+pub struct Session {
+    inner: *mut tf::TF_Session,
+}
+
+impl Session {
+    /// Creates a session.
+    pub fn new(options: &SessionOptions, graph: &Graph) -> Result<Self> {
+        let mut status = Status::new();
+        let inner = unsafe { tf::TF_NewSession(graph.inner(), options.inner, status.inner()) };
         if inner.is_null() {
             Err(status)
         } else {
             Ok(Session { inner: inner })
         }
     }
 
-    /// Loads a session from an exported model, creating a bundle
-    pub fn from_saved_model_to_bundle<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
+    /// Loads a session from an exported model.
+    pub fn from_saved_model<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
         (options: &SessionOptions,
          tags: Tags,
          graph: &mut Graph,
          export_dir: P)
-         -> Result<SavedModelBundle> {
+         -> Result<Self> {
         let mut status = Status::new();
 
         let export_dir_cstr =
@@ -102,29 +112,23 @@ impl Session {
                                          .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();
 
-        // The empty TF_Buffer will be filled by LoadSessionFromSavedModel
-        let mut meta = unsafe { Buffer::<u8>::from_ptr(ptr::null_mut(), 0) };
-
         let inner = unsafe {
             tf::TF_LoadSessionFromSavedModel(options.inner,
                                              ptr::null(),
                                              export_dir_cstr.as_ptr(),
                                              tags_ptr.as_ptr(),
                                              tags_ptr.len() as c_int,
                                              graph.inner(),
-                                             meta.inner_mut(),
+                                             ptr::null_mut(),
                                              status.inner())
         };
         if inner.is_null() {
             Err(status)
         } else {
-            let session = Session { inner: inner };
-            Ok(SavedModelBundle {
-                session: session,
-                meta_graph_def: Vec::from(meta.as_ref())
-            })
+            Ok(Session { inner: inner })
         }
     }