Closed
Description
Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template
System information
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04 x86_64): Windows 10
- TensorFlow Python installation : pip install tensorflow==2.4.1
- TensorFlow Java version : 2.3.1
- Java version (i.e., the output of
java -version): Java 8 - Python version (if transferring a model trained in Python): python 3.8
- Bazel version (if compiling from source):
- GCC/Compiler version (if compiling from source):
- CUDA/cuDNN version: No Cuda
- GPU model and memory: No GPU
Describe the current behavior
I keep getting the Error No Operation named [StatefulPartitionedCall_2:0] in the Graph
when Using SavedModelBundle.exporter to save the model
ConcreteFunction serveFunction = savedModel.function("serve_model");
SavedModelBundle.exporter(exportDir)
.withFunction(serveFunction)
.export();To access and inspect Graph operations, i can see the StatefulPartitionedCall_2
But without the : at the end of the operation name.
Iterator<Operation> operationIterator = serveFunction.graph().operations();
while(operationIterator.hasNext()){
System.out.println(operationIterator.next().name());
}code snippet output
Adam/iter
Adam/iter/Read/ReadVariableOp
Adam/beta_1
Adam/beta_1/Read/ReadVariableOp
Adam/beta_2
...
...
...
train_model_labels
StatefulPartitionedCall_1
saver_filename
StatefulPartitionedCall_2
StatefulPartitionedCall_3Works fine when invoking directly the Op from session.runner()
String checkpointPath = "...";
session.runner()
.feed("saver_filename:0", checkpointPath)
.fetch("StatefulPartitionedCall_2:0").run() ;The model was saved using a regular tf.saved_model.save function
def save_module(module, model_dir):
tf.saved_model.save(module, model_dir,
signatures={
'serve_model':
module.__call__.get_concrete_function(signature_dict),
'train_model':
module.train.get_concrete_function(
signature_dict,
tf.TensorSpec(shape=[None, None], dtype=tf.int64, name="labels"))})How can I reproduce the error:
Error could be reproduced using this scripts which defines than saves the model (credits to Thierry Herrmann)
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers
def make_model():
class CustomLayer(keras.layers.Layer):
def __init__(self, **kwargs):
super().__init__(**kwargs)
l2_reg = keras.regularizers.l2(0.1)
self.dense = layers.Dense(1, kernel_regularizer=l2_reg,
name='my_layer_dense')
def call(self, data):
return self.dense(data)
inputs = keras.Input(shape=(8,))
x1 = layers.Dense(30, activation="relu", name='my_dense')(inputs)
outputs = CustomLayer()(x1)
return keras.Model(inputs=inputs, outputs=outputs)
class CustomModule(tf.Module):
def __init__(self):
super(CustomModule, self).__init__()
self.model = make_model()
self.opt = keras.optimizers.Adam(learning_rate=0.001)
@tf.function(input_signature=[tf.TensorSpec([None, 8], tf.float32)])
def __call__(self, X):
return self.model(X)
# the my_train function processes one batch (one step): computes the loss and apply the
# loss gradient to update the model weights
@tf.function(input_signature=[tf.TensorSpec([None, 8], tf.float32), tf.TensorSpec([None], tf.float32)])
def my_train(self, X, y):
with tf.GradientTape() as tape:
logits = self.model(X, training=True)
main_loss = tf.reduce_mean(keras.losses.mean_squared_error(y, logits))
# self.model.losses contains the reularization loss (see l2_reg above)
loss_value = tf.add_n([main_loss] + self.model.losses)
grads = tape.gradient(loss_value, self.model.trainable_weights)
self.opt.apply_gradients(zip(grads, self.model.trainable_weights))
return loss_value
# instantiate the module
module = CustomModule()
def save_module(module, model_dir):
tf.saved_model.save(module, model_dir,
signatures={
'serve_model' :
module.__call__.get_concrete_function(tf.TensorSpec([None, 8], tf.float32)),
'train_model' :
module.my_train.get_concrete_function(tf.TensorSpec([None, 8], tf.float32),
tf.TensorSpec([None], tf.float32))})
MODEL_OUTPUT_DIR ="..."
save_module(module, MODEL_OUTPUT_DIR)@karllessard could you please take a quick look ? Thank you
Metadata
Metadata
Assignees
Labels
No labels