Merge pull request #159 from bekker/feature/example-checkpoint

Add an example for saving and loading checkpoints

Author: adamcrume

Cargo.toml

@@ -38,3 +38,6 @@ name = "regression"
 
 [[example]]
 name = "regression_savedmodel"
+
+[[example]]
+name = "regression_checkpoint"

examples/addition.rs

@@ -32,7 +32,7 @@ fn main() {
 }
 
 fn run() -> Result<(), Box<Error>> {
-    let filename = "examples/addition-model/model.pb"; // z = x + y
+    let filename = "examples/addition/model.pb"; // z = x + y
     if !Path::new(filename).exists() {
         return Err(Box::new(Status::new_set(Code::NotFound,
                                             &format!("Run 'python addition.py' to generate {} \

examples/addition.py renamed to examples/addition/addition.py

@@ -8,5 +8,5 @@
 tf.variables_initializer(tf.global_variables(), name = 'init')
 
 definition = tf.Session().graph_def
-directory = 'examples/addition-model'
+directory = 'examples/addition'
 tf.train.write_graph(definition, directory, 'model.pb', as_text=False)

examples/addition/model.pb

@@ -0,0 +1,14 @@
+
+,
+
xPlaceholder*
+shape:
*
+dtype0
+,
+
yPlaceholder*
+dtype0*
+shape:

+
+
zAdd
x
y*
+
T0
+
+initNoOp"

examples/regression.rs

@@ -34,7 +34,7 @@ fn main() {
 }
 
 fn run() -> Result<(), Box<Error>> {
-    let filename = "examples/regression-model/model.pb"; // y = w * x + b
+    let filename = "examples/regression/model.pb"; // y = w * x + b
     if !Path::new(filename).exists() {
         return Err(Box::new(Status::new_set(Code::NotFound,
                                             &format!("Run 'python regression.py' to generate \
@@ -71,8 +71,6 @@ fn run() -> Result<(), Box<Error>> {
 
     // Load the test data into the session.
     let mut init_step = StepWithGraph::new();
-    init_step.add_input(&op_x, 0, &x);
-    init_step.add_input(&op_y, 0, &y);
     init_step.add_target(&op_init);
     session.run(&mut init_step)?;
 

examples/regression-model/model.pb renamed to examples/regression/model.pb

@@ -15,5 +15,5 @@
 init = tf.variables_initializer(tf.global_variables(), name='init')
 
 definition = tf.Session().graph_def
-directory = 'examples/regression-model'
+directory = 'examples/regression'
 tf.train.write_graph(definition, directory, 'model.pb', as_text=False)

examples/regression.py renamed to examples/regression/regression.py

@@ -0,0 +1,131 @@
+#![cfg_attr(feature="nightly", feature(alloc_system))]
+#[cfg(feature="nightly")]
+extern crate alloc_system;
+extern crate random;
+extern crate tensorflow;
+
+use random::Source;
+use std::error::Error;
+use std::fs::File;
+use std::io::Read;
+use std::result::Result;
+use std::path::Path;
+use std::process::exit;
+use tensorflow::Code;
+use tensorflow::Graph;
+use tensorflow::ImportGraphDefOptions;
+use tensorflow::Session;
+use tensorflow::SessionOptions;
+use tensorflow::Status;
+use tensorflow::StepWithGraph;
+use tensorflow::Tensor;
+
+fn main() {
+    // Putting the main code in another function serves two purposes:
+    // 1. We can use the `?` operator.
+    // 2. We can call exit safely, which does not run any destructors.
+    exit(match run() {
+        Ok(_) => 0,
+        Err(e) => {
+            println!("{}", e);
+            1
+        }
+    })
+}
+
+fn run() -> Result<(), Box<Error>> {
+    let filename = "examples/regression_checkpoint/model.pb"; // y = w * x + b
+    if !Path::new(filename).exists() {
+        return Err(Box::new(Status::new_set(Code::NotFound,
+                                            &format!("Run 'python regression_checkpoint.py' to generate \
+                                                      {} and try again.",
+                                                     filename))
+            .unwrap()));
+    }
+
+    // Generate some test data.
+    let w = 0.1;
+    let b = 0.3;
+    let num_points = 100;
+    let steps = 201;
+    let mut rand = random::default();
+    let mut x = Tensor::new(&[num_points as u64]);
+    let mut y = Tensor::new(&[num_points as u64]);
+    for i in 0..num_points {
+        x[i] = (2.0 * rand.read::<f64>() - 1.0) as f32;
+        y[i] = w * x[i] + b;
+    }
+
+    // Load the computation graph defined by regression.py.
+    let mut graph = Graph::new();
+    let mut proto = Vec::new();
+    File::open(filename)?.read_to_end(&mut proto)?;
+    graph.import_graph_def(&proto, &ImportGraphDefOptions::new())?;
+    let mut session = Session::new(&SessionOptions::new(), &graph)?;
+    let op_x = graph.operation_by_name_required("x")?;
+    let op_y = graph.operation_by_name_required("y")?;
+    let op_init = graph.operation_by_name_required("init")?;
+    let op_train = graph.operation_by_name_required("train")?;
+    let op_w = graph.operation_by_name_required("w")?;
+    let op_b = graph.operation_by_name_required("b")?;
+    let op_file_path = graph.operation_by_name_required("save/Const")?;
+    let op_save = graph.operation_by_name_required("save/control_dependency")?;
+    let file_path_tensor: Tensor<String> = Tensor::from(String::from("examples/regression_checkpoint/saved.ckpt"));
+
+    // Load the test data into the session.
+    let mut init_step = StepWithGraph::new();
+    init_step.add_target(&op_init);
+    session.run(&mut init_step)?;
+
+    // Train the model.
+    let mut train_step = StepWithGraph::new();
+    train_step.add_input(&op_x, 0, &x);
+    train_step.add_input(&op_y, 0, &y);
+    train_step.add_target(&op_train);
+    for _ in 0..steps {
+        session.run(&mut train_step)?;
+    }
+
+    // Save the model.
+    let mut step = StepWithGraph::new();
+    step.add_input(&op_file_path, 0, &file_path_tensor);
+    step.add_target(&op_save);
+    session.run(&mut step)?;
+
+    // Initialize variables, to erase trained data.
+    session.run(&mut init_step)?;
+
+    // Load the model.
+    let op_load = graph.operation_by_name_required("save/restore_all")?;
+    let mut step = StepWithGraph::new();
+    step.add_input(&op_file_path, 0, &file_path_tensor);
+    step.add_target(&op_load);
+    session.run(&mut step)?;
+
+    // Grab the data out of the session.
+    let mut output_step = StepWithGraph::new();
+    let w_ix = output_step.request_output(&op_w, 0);
+    let b_ix = output_step.request_output(&op_b, 0);
+    session.run(&mut output_step)?;
+
+    // Check our results.
+    let w_hat: f32 = output_step.take_output(w_ix)?[0];
+    let b_hat: f32 = output_step.take_output(b_ix)?[0];
+    println!("Checking w: expected {}, got {}. {}",
+             w,
+             w_hat,
+             if (w - w_hat).abs() < 1e-3 {
+                 "Success!"
+             } else {
+                 "FAIL"
+             });
+    println!("Checking b: expected {}, got {}. {}",
+             b,
+             b_hat,
+             if (b - b_hat).abs() < 1e-3 {
+                 "Success!"
+             } else {
+                 "FAIL"
+             });
+    Ok(())
+}

examples/regression_checkpoint.rs

@@ -0,0 +1 @@
+saved.ckpt*

examples/regression_checkpoint/.gitignore

@@ -0,0 +1,22 @@
+import os
+import tensorflow as tf
+
+x = tf.placeholder(tf.float32, name='x')
+y = tf.placeholder(tf.float32, name='y')
+
+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
+
+loss = tf.reduce_mean(tf.square(y_hat - y))
+optimizer = tf.train.GradientDescentOptimizer(0.5)
+train = optimizer.minimize(loss, name='train')
+
+init = tf.variables_initializer(tf.global_variables(), name='init')
+
+# Declare saver ops
+saver = tf.train.Saver(tf.global_variables())
+
+definition = tf.Session().graph_def
+directory = 'examples/regression_checkpoint'
+tf.train.write_graph(definition, directory, 'model.pb', as_text=False)

examples/regression_checkpoint/model.pb

@@ -31,7 +31,7 @@ fn main() {
 }
 
 fn run() -> Result<(), Box<Error>> {
-    let export_dir = "examples/saved-regression-model"; // y = w * x + b
+    let export_dir = "examples/regression_savedmodel"; // y = w * x + b
     if !Path::new(export_dir).exists() {
         return Err(Box::new(Status::new_set(Code::NotFound,
                                             &format!("Run 'python regression_savedmodel.py' to generate \

examples/regression_checkpoint/regression_checkpoint.py

@@ -18,7 +18,7 @@
 
 init = tf.variables_initializer(tf.global_variables(), name='init')
 
-directory = 'examples/saved-regression-model'
+directory = 'examples/regression_savedmodel'
 builder = SavedModelBuilder(directory)
 
 with tf.Session(graph=tf.get_default_graph()) as sess: