Skip to content

Update regression_savedmodel example for tf2.5 #312

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 11, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 25 additions & 12 deletions examples/regression_savedmodel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,18 +46,31 @@ 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(),
&["train", "serve"],
&mut graph,
export_dir,
)?
.session;
let op_x = graph.operation_by_name_required("x")?;
let op_y = graph.operation_by_name_required("y")?;
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 bundle =
SavedModelBundle::load(&SessionOptions::new(), &["serve"], &mut graph, export_dir)?;
let session = &bundle.session;

// train
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 loss_info = train_signature.get_output("loss")?;
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(&loss_info.name().name)?;

// internal parameters
let op_b = {
let b_signature = bundle.meta_graph_def().get_signature("b")?;
let b_info = b_signature.get_output("output")?;
graph.operation_by_name_required(&b_info.name().name)?
};

let op_w = {
let w_signature = bundle.meta_graph_def().get_signature("w")?;
let w_info = w_signature.get_output("output")?;
graph.operation_by_name_required(&w_info.name().name)?
};

// Train the model (e.g. for fine tuning).
let mut train_step = SessionRunArgs::new();
Expand Down
Empty file.
76 changes: 40 additions & 36 deletions examples/regression_savedmodel/regression_savedmodel.py
Original file line number Diff line number Diff line change
@@ -1,43 +1,47 @@
import tensorflow as tf
from tensorflow.python.saved_model.builder import SavedModelBuilder
from tensorflow.python.saved_model.signature_def_utils import build_signature_def
from tensorflow.python.saved_model.signature_constants import REGRESS_METHOD_NAME
from tensorflow.python.saved_model.tag_constants import TRAINING, SERVING
from tensorflow.python.saved_model.utils import build_tensor_info

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 = tf.add(w * x, b, name="y_hat")
class LinearRegresstion(tf.Module):
def __init__(self, name=None):
super(LinearRegresstion, self).__init__(name=name)
self.w = tf.Variable(tf.random.uniform([1], -1.0, 1.0), name='w')
self.b = tf.Variable(tf.zeros([1]), name='b')
self.optimizer = tf.keras.optimizers.SGD(0.5)

loss = tf.reduce_mean(tf.square(y_hat - y))
optimizer = tf.train.GradientDescentOptimizer(0.5)
train = optimizer.minimize(loss, name='train')
@tf.function
def __call__(self, x):
y_hat = self.w * x + self.b
return y_hat

init = tf.variables_initializer(tf.global_variables(), name='init')
@tf.function
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))
return {'loss': loss}


model = LinearRegresstion()

# Get concrete functions to generate signatures
x = tf.TensorSpec([None], tf.float32, name='x')
y = tf.TensorSpec([None], tf.float32, name='y')

train = model.train.get_concrete_function(x, y)
w = model.get_w.get_concrete_function()
b = model.get_b.get_concrete_function()

signatures = {'train': train, 'w': w, 'b': b}

directory = 'examples/regression_savedmodel'
builder = SavedModelBuilder(directory)

with tf.Session(graph=tf.get_default_graph()) as sess:
sess.run(init)

signature_inputs = {
"x": build_tensor_info(x),
"y": build_tensor_info(y)
}
signature_outputs = {
"out": build_tensor_info(y_hat)
}
signature_def = build_signature_def(
signature_inputs, signature_outputs,
REGRESS_METHOD_NAME)
builder.add_meta_graph_and_variables(
sess, [TRAINING, SERVING],
signature_def_map={
REGRESS_METHOD_NAME: signature_def
},
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS))
builder.save(as_text=False)
tf.saved_model.save(model, directory, signatures=signatures)
Binary file modified examples/regression_savedmodel/saved_model.pb
Binary file not shown.
Binary file not shown.
Binary file modified examples/regression_savedmodel/variables/variables.index
Binary file not shown.