Merge pull request tensorflow#232 from adamcrume/saved-model

Add support for creating saved models

Author: adamcrume

Cargo.toml

@@ -18,6 +18,8 @@ tensorflow-sys = { version = "0.17.0", path = "tensorflow-sys" }
 byteorder = "1.3.2"
 crc = "1.8.1"
 half = "1.3.0"
+# This is used internally but not intended to be exposed through the API.
+protobuf = "=2.8.0"
 
 [dev-dependencies]
 random = "0.12.2"

RELEASING.md

@@ -1,3 +1,10 @@
+## Pre-release
+
+1. Generate protos
+   1. Run `cd tensorflow-proto-codegen; cargo run -- $PATH_TO_TENSORFLOW $PWD/../src/protos`
+   1. Update Cargo.toml to ensure version of protobuf exactly equals version of protoc_rust used
+1. Commit and push changes
+
 ## Releasing
 
 1. Check out a clean copy.  Note that `cargo publish` packages up untracked files.  Use `--allow-dirty` at your peril.

examples/xor.rs

@@ -1,20 +1,32 @@
+use std::env;
 use std::error::Error;
+use std::fs;
+use std::io::ErrorKind;
+use std::path::Path;
 use std::result::Result;
 use tensorflow::ops;
 use tensorflow::train::AdadeltaOptimizer;
 use tensorflow::train::MinimizeOptions;
 use tensorflow::train::Optimizer;
 use tensorflow::Code;
 use tensorflow::DataType;
+use tensorflow::Graph;
 use tensorflow::Output;
+use tensorflow::OutputName;
+use tensorflow::SavedModelBundle;
 use tensorflow::Scope;
 use tensorflow::Session;
 use tensorflow::SessionOptions;
 use tensorflow::SessionRunArgs;
 use tensorflow::Shape;
+use tensorflow::SignatureDef;
 use tensorflow::Status;
 use tensorflow::Tensor;
+use tensorflow::TensorInfo;
 use tensorflow::Variable;
+use tensorflow::REGRESS_INPUTS;
+use tensorflow::REGRESS_METHOD_NAME;
+use tensorflow::REGRESS_OUTPUTS;
 
 // Helper for building a layer.
 //
@@ -54,7 +66,7 @@ fn layer<O1: Into<Output>>(
     ))
 }
 
-fn main() -> Result<(), Box<Error>> {
+fn train<P: AsRef<Path>>(save_dir: P) -> Result<(), Box<dyn Error>> {
     // ================
     // Build the model.
     // ================
@@ -94,6 +106,34 @@ fn main() -> Result<(), Box<Error>> {
         MinimizeOptions::default().with_variables(&variables),
     )?;
 
+    let mut all_vars = variables.clone();
+    all_vars.extend_from_slice(&minimizer_vars);
+    let mut builder = tensorflow::SavedModelBuilder::new();
+    builder
+        .add_collection("train", &all_vars)
+        .add_tag("serve")
+        .add_tag("train")
+        .add_signature(REGRESS_METHOD_NAME, {
+            let mut def = SignatureDef::new(REGRESS_METHOD_NAME.to_string());
+            def.add_input_info(
+                REGRESS_INPUTS.to_string(),
+                TensorInfo::new(
+                    DataType::Float,
+                    Shape::from(None),
+                    OutputName {
+                        name: input.name()?,
+                        index: 0,
+                    },
+                ),
+            );
+            def.add_output_info(
+                REGRESS_OUTPUTS.to_string(),
+                TensorInfo::new(DataType::Float, Shape::from(None), layer2.name()?),
+            );
+            def
+        });
+    let saved_model_saver = builder.inject(scope)?;
+
     // =========================
     // Initialize the variables.
     // =========================
@@ -118,7 +158,7 @@ fn main() -> Result<(), Box<Error>> {
     let mut label_tensor = Tensor::<f32>::new(&[1]);
     // Helper that generates a training example from an integer, trains on that
     // example, and returns the error.
-    let mut train = |i| -> Result<f32, Box<Error>> {
+    let mut train = |i| -> Result<f32, Box<dyn Error>> {
         input_tensor[0] = (i & 1) as f32;
         input_tensor[1] = ((i >> 1) & 1) as f32;
         label_tensor[0] = ((i & 1) ^ ((i >> 1) & 1)) as f32;
@@ -134,6 +174,11 @@ fn main() -> Result<(), Box<Error>> {
         train(i)?;
     }
 
+    // ================
+    // Save the model.
+    // ================
+    saved_model_saver.save(&session, &g, &save_dir)?;
+
     // ===================
     // Evaluate the model.
     // ===================
@@ -149,3 +194,70 @@ fn main() -> Result<(), Box<Error>> {
     }
     Ok(())
 }
+
+fn eval<P: AsRef<Path>>(save_dir: P) -> Result<(), Box<dyn Error>> {
+    let mut graph = Graph::new();
+    let bundle = SavedModelBundle::load(
+        &SessionOptions::new(),
+        &["serve", "train"],
+        &mut graph,
+        save_dir,
+    )?;
+    let session = &bundle.session;
+    let signature = bundle.meta_graph_def().get_signature(REGRESS_METHOD_NAME)?;
+    let input_info = signature.get_input(REGRESS_INPUTS)?;
+    let output_info = signature.get_output(REGRESS_OUTPUTS)?;
+    let input_op = graph.operation_by_name_required(&input_info.name().name)?;
+    let output_op = graph.operation_by_name_required(&output_info.name().name)?;
+
+    let mut input_tensor = Tensor::<f32>::new(&[1, 2]);
+    for i in 0..4 {
+        input_tensor[0] = (i & 1) as f32;
+        input_tensor[1] = ((i >> 1) & 1) as f32;
+        let expected = ((i & 1) ^ ((i >> 1) & 1)) as f32;
+        let mut run_args = SessionRunArgs::new();
+        run_args.add_feed(&input_op, input_info.name().index, &input_tensor);
+        let output_fetch = run_args.request_fetch(&output_op, output_info.name().index);
+        session.run(&mut run_args)?;
+        let output = run_args.fetch::<f32>(output_fetch)?[0];
+        let error = (output - expected) * (output - expected);
+        println!("Error: {}", error);
+        if error > 0.1 {
+            return Err(Box::new(Status::new_set(
+                Code::Internal,
+                &format!("Error too high: {}", error),
+            )?));
+        }
+    }
+
+    Ok(())
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let mut dir = env::temp_dir();
+    dir.push("tf-rust-example-xor-saved-model");
+    let mut dir2 = env::temp_dir();
+    dir2.push("tf-rust-example-xor-saved-model2");
+    match fs::remove_dir_all(&dir) {
+        Err(e) => {
+            if e.kind() != ErrorKind::NotFound {
+                return Err(Box::new(e));
+            }
+        }
+        Ok(_) => (),
+    }
+    match fs::remove_dir_all(&dir2) {
+        Err(e) => {
+            if e.kind() != ErrorKind::NotFound {
+                return Err(Box::new(e));
+            }
+        }
+        Ok(_) => (),
+    }
+    train(&dir)?;
+    // Ensure that the saved model works even when moved.
+    // Users do not need to do this; this is purely for testing purposes.
+    fs::rename(&dir, &dir2)?;
+    eval(&dir2)?;
+    Ok(())
+}

src/graph.rs

@@ -18,9 +18,13 @@ use std;
 use std::ffi::CStr;
 use std::ffi::CString;
 use std::ffi::NulError;
+use std::fmt;
+use std::fmt::Display;
+use std::fmt::Formatter;
 use std::os::raw::c_void as std_c_void;
 use std::ptr;
 use std::slice;
+use std::str::FromStr;
 use std::str::Utf8Error;
 use std::sync::Arc;
 use tensorflow_sys as tf;
@@ -1724,6 +1728,50 @@ impl Output {
             })
         }
     }
+
+    /// Returns the name of this output.
+    pub fn name(&self) -> Result<OutputName> {
+        Ok(OutputName {
+            name: self.operation.name()?,
+            index: self.index,
+        })
+    }
+}
+
+////////////////////////
+
+/// Names a specific Output in the graph.
+#[derive(Clone, PartialEq, Eq, Hash, Debug, Default)]
+pub struct OutputName {
+    /// Name of the operation the edge connects to.
+    pub name: String,
+
+    /// Index into either the outputs of the operation.
+    pub index: c_int,
+}
+
+impl FromStr for OutputName {
+    type Err = Status;
+    fn from_str(s: &str) -> Result<Self> {
+        let splits: Vec<_> = s.split(':').collect();
+        if splits.len() != 2 {
+            return Err(Status::new_set_lossy(
+                Code::InvalidArgument,
+                "Name must contain exactly one colon (':')",
+            ));
+        }
+        let index = splits[1].parse::<c_int>()?;
+        Ok(Self {
+            name: splits[0].to_string(),
+            index,
+        })
+    }
+}
+
+impl Display for OutputName {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.name, self.index)
+    }
 }
 
 ////////////////////////
@@ -2847,4 +2895,26 @@ mod tests {
         assert_eq!(consumers[0].0.name().unwrap(), "y");
         assert_eq!(consumers[0].1, 0);
     }
+
+    #[test]
+    fn output_name() {
+        assert_eq!(
+            "foo:1".parse::<OutputName>().unwrap(),
+            OutputName {
+                name: "foo".to_string(),
+                index: 1
+            }
+        );
+        assert_eq!(
+            OutputName {
+                name: "foo".to_string(),
+                index: 1
+            }
+            .to_string(),
+            "foo:1"
+        );
+        assert!("foo".parse::<OutputName>().is_err());
+        assert!("foo:bar".parse::<OutputName>().is_err());
+        assert!("foo:0:1".parse::<OutputName>().is_err());
+    }
 }