opの名前の取得方法をsignature経由に変更する

Author: dskkato

examples/regression_savedmodel.rs

@@ -41,27 +41,41 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // Load the saved model exported by regression_savedmodel.py.
     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("train_x")?;
-    let op_y = graph.operation_by_name_required("train_y")?;
-    let op_train = graph.operation_by_name_required("StatefulPartitionedCall")?;
-    let op_w = graph.operation_by_name_required("StatefulPartitionedCall_1")?;
-    let op_b = graph.operation_by_name_required("StatefulPartitionedCall_1")?;
+    let bundle =
+        SavedModelBundle::load(&SessionOptions::new(), &["serve"], &mut graph, export_dir)?;
+    let session = &bundle.session;
+
+    let train_signature = bundle.meta_graph_def().get_signature("train")?;
+    let x_info = train_signature.get_input("x")?;
+    let y_info = train_signature.get_input("y")?;
+    let train_info = train_signature.get_output("train")?;
+    let op_x = graph.operation_by_name_required(&x_info.name().name)?;
+    let op_y = graph.operation_by_name_required(&y_info.name().name)?;
+    let op_train = graph.operation_by_name_required(&train_info.name().name)?;
+    let w_info = bundle
+        .meta_graph_def()
+        .get_signature("w")?
+        .get_output("output")?;
+    let op_w = graph.operation_by_name_required(&w_info.name().name)?;
+    let b_info = bundle
+        .meta_graph_def()
+        .get_signature("b")?
+        .get_output("output")?;
+    let op_b = graph.operation_by_name_required(&b_info.name().name)?;
 
     // Train the model (e.g. for fine tuning).
     let mut train_step = SessionRunArgs::new();
     train_step.add_feed(&op_x, 0, &x);
     train_step.add_feed(&op_y, 0, &y);
-    train_step.request_fetch(&op_train, 0);
+    train_step.add_target(&op_train);
     for _ in 0..steps {
         session.run(&mut train_step)?;
     }
 
     // Grab the data out of the session.
     let mut output_step = SessionRunArgs::new();
     let w_ix = output_step.request_fetch(&op_w, 0);
-    let b_ix = output_step.request_fetch(&op_b, 1);
+    let b_ix = output_step.request_fetch(&op_b, 0);
     session.run(&mut output_step)?;
 
     // Check our results.

examples/regression_savedmodel/regression_savedmodel.py

@@ -14,35 +14,39 @@ def __call__(self, x):
         return y_hat
 
     @tf.function
-    def get_weights(self):
-        return self.w, self.b
+    def get_w(self):
+        return {"output": self.w}
+
+    @tf.function
+    def get_b(self):
+        return {"output": self.b}
 
     @tf.function
     def train(self, x, y):
         with tf.GradientTape() as tape:
             y_hat = self(x)
             loss = tf.reduce_mean(tf.square(y_hat - y))
         grads = tape.gradient(loss, self.trainable_variables)
-        _ = self.optimizer.apply_gradients(
-            zip(grads, self.trainable_variables), name="train"
-        )
-        return loss
+        _ = self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
+        return {"train": loss}
 
 
 model = LinearRegresstion()
 
 x = tf.TensorSpec([None], tf.float32, name="x")
 y = tf.TensorSpec([None], tf.float32, name="y")
 train = model.train.get_concrete_function(x, y)
-weights = model.get_weights.get_concrete_function()
+w = model.get_w.get_concrete_function()
+b = model.get_b.get_concrete_function()
 
 directory = "examples/regression_savedmodel"
-signatures = {"train": train, "weights": weights}
+signatures = {"train": train, "w": w, "b": b}
 tf.saved_model.save(model, directory, signatures=signatures)
 
 # export graph info to TensorBoard
 logdir = "logs/regression_savedmodel"
 writer = tf.summary.create_file_writer(logdir)
 with writer.as_default():
     tf.summary.graph(train.graph)
-    tf.summary.graph(weights.graph)
+    tf.summary.graph(w.graph)
+    tf.summary.graph(b.graph)