bhack
commented
3 years ago Open klaimans opened 3 years ago
bhack
commented
3 years ago 
klaimans
commented
3 years ago Hi @bhack , Thank you very much for the very quick reply. Indeed that gets us past that error but if you try the same for the decoder rather than the encoder (which is what we are trying to do) it still fails. It is also unclear how to send the signature for the AttenstionWrapperState input to the decoder which is a NamedTuple and not a simple list.
I would appreciate any additional suggestions you might have. What we see in our subclassed model is that while we can create a concrete function for our predict step the save_model function still fails since it tries to trace the model with the wrong shaped tensors even though for our specific predict_step there is no problem and we define the signatures explicitly.
Thank you in advance!
guillaumekln
commented
3 years ago Did you manage to solve the issue? Did you try saving the decoder in a checkpoint instead?
If you need to save a self-contained graph, I think you should define a separate tf.function wrapping the decoder with a simpler signature. Then you could export this function with tf.saved_model.save.
klaimans
commented
3 years ago Hi @guillaumekln , sorry for the belated reply. Unfortunately we are still struggling with this. Using a checkpoint works without a problem and as I mentioned we can even create a concrete function which works perfectly. However, sending the model to tf.save_model.save still fails. We have tried different signatures also overriding some of the builtin signatures in tfa.Sampler and tfa.BaseDecoder but we couldn't resolve it yet.
We ended up solving this by using the logic for saving estimator based models. It is ugly but at least we got it to work and we could load the model and use it for inference.
We would really appreciate any additional suggestions.
guillaumekln
commented
3 years ago Here's an example that exports a BasicDecoder to a SavedModel:
import tempfile
import tensorflow as tf
import tensorflow_addons as tfa
class MyModel(tf.keras.layers.Layer):
def __init__(self, vocab_size, num_units):
super().__init__()
self.embedding = tf.keras.layers.Embedding(vocab_size, num_units)
self.cell = tf.keras.layers.LSTMCell(num_units)
self.sampler = tfa.seq2seq.GreedyEmbeddingSampler(self.embedding)
self.output_layer = tf.keras.layers.Dense(vocab_size)
self.decoder = tfa.seq2seq.BasicDecoder(
self.cell,
self.sampler,
self.output_layer,
maximum_iterations=10,
)
@tf.function(
input_signature=(
tf.TensorSpec([None], dtype=tf.int32),
tf.TensorSpec([], dtype=tf.int32),
)
)
def run(self, start_tokens, end_token):
batch_size = tf.shape(start_tokens)[0]
initial_state = self.cell.get_initial_state(
batch_size=batch_size, dtype=tf.float32
)
output, state, lengths = self.decoder(
None,
start_tokens=start_tokens,
end_token=end_token,
initial_state=initial_state,
)
return output
with tempfile.TemporaryDirectory() as export_dir:
model = MyModel(512, 64)
tf.saved_model.save(model, export_dir, signatures=model.run.get_concrete_function())
del model
imported = tf.saved_model.load(export_dir)
function = imported.signatures["serving_default"]
start_tokens = tf.constant([1, 2, 3], dtype=tf.int32)
end_token = tf.constant(5, dtype=tf.int32)
output = function(start_tokens=start_tokens, end_token=end_token)
print(output)
System information
Describe the bug
when trying to save a subclassed model with a BasicDecoder layer, one cannot save the model using the save method or tf.saved_model.save
Code to reproduce the issue
Using the colab tutorial
https://colab.research.google.com/github/tensorflow/addons/blob/master/docs/tutorials/networks_seq2seq_nmt.ipynb
if one tries to save the decoder it doesn't work.
Our subclassed model is more involved but we expect that if we cannot save the decoder already in this example it cannot work for us either.
Provide a reproducible test case that is the bare minimum necessary to generate the problem.
see colab above and add after the training cell the following line:
decoder.save("./test")
TypeError Traceback (most recent call last)