Skip to content

Don't save traces for saved model #945

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.

Sign up for GitHub

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
Apr 1, 2023
Merged

Don't save traces for saved model #945

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 3 additions & 1 deletion keras_nlp/layers/f_net_encoder_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,9 @@ def test_saved_model(self, save_format, filename):
data = tf.random.uniform(shape=[2, 4, 6])
model(data)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
loaded_model = keras.models.load_model(path)

model_output = model(data)
Expand Down
19 changes: 13 additions & 6 deletions keras_nlp/layers/masked_lm_head_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,13 @@
import os

import tensorflow as tf
from absl.testing import parameterized
from tensorflow import keras

from keras_nlp.layers import masked_lm_head


class MaskedLMHeadTest(tf.test.TestCase):
class MaskedLMHeadTest(tf.test.TestCase, parameterized.TestCase):
def test_valid_call(self):
head = masked_lm_head.MaskedLMHead(
vocabulary_size=100,
Expand Down Expand Up @@ -156,7 +157,11 @@ def test_checkpointing(self):
head2_output = head2(token_data, mask_positions=position_data)
self.assertAllClose(head1_output, head2_output)

def test_saving_model(self):
@parameterized.named_parameters(
("tf_format", "tf", "model"),
("keras_format", "keras_v3", "model.keras"),
)
def test_saved_model(self, save_format, filename):
head = masked_lm_head.MaskedLMHead(
vocabulary_size=100,
activation="softmax",
Expand All @@ -171,9 +176,11 @@ def test_saving_model(self):
shape=(4, 5), maxval=10, dtype="int32"
)
model_output = model((token_data, position_data))
save_path = os.path.join(self.get_temp_dir(), "model")
model.save(save_path)
restored = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

restored_output = restored((token_data, position_data))
restored_output = restored_model((token_data, position_data))
self.assertAllClose(model_output, restored_output)
4 changes: 3 additions & 1 deletion keras_nlp/layers/multi_segment_packer_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,9 @@ def test_saved_model(self, save_format, filename):
outputs = packer(inputs)
model = keras.Model(inputs, outputs)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)
self.assertAllEqual(
model((seq1, seq2)),
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/layers/position_embedding_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,7 +309,9 @@ def test_saved_model(self, save_format, filename):
model(data)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
loaded_model = keras.models.load_model(path)

model_output = model.predict(data)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/layers/token_and_position_embedding_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,9 @@ def test_saved_model(self, save_format, filename):
model(data)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
loaded_model = keras.models.load_model(path)

model_output = model.predict(data)
Expand Down
8 changes: 6 additions & 2 deletions keras_nlp/layers/transformer_decoder_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,7 +332,9 @@ def test_saved_model(self, save_format, filename):
decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
model([decoder_sequence, encoder_sequence])
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

loaded_model = keras.models.load_model(path)
model_output = model([decoder_sequence, encoder_sequence])
Expand All @@ -358,7 +360,9 @@ def test_saved_model_without_cross_attention(self, save_format, filename):
decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
model(decoder_sequence)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
loaded_model = keras.models.load_model(path)

model_output = model(decoder_sequence)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/layers/transformer_encoder_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,9 @@ def test_saved_model(self, save_format, filename):
data = tf.random.uniform(shape=[2, 4, 6])
model_output = model(data)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

loaded_model = keras.models.load_model(path)
loaded_model_output = loaded_model(data)
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/albert/albert_backbone_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,9 +93,11 @@ def test_error_for_invalid_num_groups(self):
@pytest.mark.large
def test_saved_model(self, save_format, filename):
model_output = self.backbone(self.input_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.backbone.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.backbone.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, AlbertBackbone)
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/albert/albert_classifier_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,9 +125,11 @@ def test_serialization(self):
@pytest.mark.large
def test_saving_model(self, save_format, filename):
model_output = self.classifier.predict(self.raw_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.classifier.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.classifier.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back
self.assertIsInstance(restored_model, AlbertClassifier)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,9 @@ def test_saved_model(self, save_format, filename):
model = keras.Model(inputs, outputs)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

restored_model = keras.models.load_model(path)
outputs = model(input_data)[0]["token_ids"]
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/albert/albert_masked_lm_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,9 +130,11 @@ def test_classifier_fit_no_xla(self):
@pytest.mark.large
def test_saved_model(self, save_format, filename):
model_output = self.masked_lm.predict(self.raw_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.masked_lm.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.masked_lm.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, AlbertMaskedLM)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/albert/albert_preprocessor_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,9 @@ def test_saved_model(self, save_format, filename):
outputs = self.preprocessor(inputs)
model = keras.Model(inputs, outputs)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)
self.assertAllEqual(
model(input_data)["token_ids"],
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/albert/albert_tokenizer_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,9 @@ def test_saved_model(self, save_format, filename):
model = keras.Model(inputs, outputs)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

restored_model = keras.models.load_model(path)
self.assertAllEqual(
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/bart/bart_backbone_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,9 +97,11 @@ def test_compile_batched_ds(self, jit_compile):
)
def test_saved_model(self, save_format, filename):
model_output = self.model(self.input_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.model.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, BartBackbone)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/bart/bart_tokenizer_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,9 @@ def test_saved_model(self, save_format, filename):
model = keras.Model(inputs, outputs)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

restored_model = keras.models.load_model(path)
self.assertAllEqual(
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/bert/bert_backbone_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,11 @@ def test_serialization(self):
@pytest.mark.large
def test_saved_model(self, save_format, filename):
model_output = self.backbone(self.input_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.backbone.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.backbone.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, BertBackbone)
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/bert/bert_classifier_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,9 +98,11 @@ def test_serialization(self):
@pytest.mark.large # Saving is slow, so mark these large.
def test_saved_model(self, save_format, filename):
model_output = self.classifier.predict(self.raw_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.classifier.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.classifier.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, BertClassifier)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/bert/bert_masked_lm_preprocessor_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,9 @@ def test_saved_model(self, save_format, filename):
model = keras.Model(inputs, outputs)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

restored_model = keras.models.load_model(path)
outputs = model(input_data)[0]["token_ids"]
Expand Down
10 changes: 6 additions & 4 deletions keras_nlp/models/bert/bert_masked_lm_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,15 @@ def test_serialization(self):
@pytest.mark.large # Saving is slow, so mark these large.
def test_saved_model(self, save_format, filename):
model_output = self.masked_lm.predict(self.raw_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.masked_lm.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.masked_lm.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, BertMaskedLM)

# Check that output matches.
restored_output = restored_model.predict(self.raw_batch)
self.assertAllClose(model_output, restored_output)
self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Curious - why would it fluctuate? 0.01 without context is not a common scale in atol.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not really sure, also this is totally unrelated, so I can split into a separate PR. I was seeing enough fluctuation on my nvidia GPU that sometimes these tests would fail. This is true on master too.

It looks like it is just a precision issue, the floats are the same, with a lax enough tolerance.

4 changes: 3 additions & 1 deletion keras_nlp/models/bert/bert_preprocessor_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,9 @@ def test_saved_model(self, save_format, filename):
outputs = self.preprocessor(inputs)
model = keras.Model(inputs, outputs)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)
self.assertAllEqual(
model(input_data)["token_ids"],
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/bert/bert_tokenizer_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,9 @@ def test_saved_model(self, save_format, filename):
outputs = tokenizer(inputs)
model = keras.Model(inputs, outputs)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)
self.assertAllEqual(
model(input_data),
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,9 +83,11 @@ def test_serialization(self):
@pytest.mark.large
def test_saved_model(self, save_format, filename):
model_output = self.backbone(self.input_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.backbone.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.backbone.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, DebertaV3Backbone)
Expand Down
8 changes: 5 additions & 3 deletions keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,11 @@ def test_serialization(self):
@pytest.mark.large
def test_saving_model(self, save_format, filename):
model_output = self.classifier.predict(self.raw_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.classifier.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.classifier.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, DebertaV3Classifier)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,9 @@ def test_saved_model(self, save_format, filename):
model = keras.Model(inputs, outputs)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

restored_model = keras.models.load_model(path)
outputs = model(input_data)[0]["token_ids"]
Expand Down
10 changes: 6 additions & 4 deletions keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,13 +118,15 @@ def test_serialization(self):
@pytest.mark.large
def test_saved_model(self, save_format, filename):
model_output = self.masked_lm.predict(self.raw_batch)
save_path = os.path.join(self.get_temp_dir(), filename)
self.masked_lm.save(save_path, save_format=save_format)
restored_model = keras.models.load_model(save_path)
path = os.path.join(self.get_temp_dir(), filename)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
self.masked_lm.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)

# Check we got the real object back.
self.assertIsInstance(restored_model, DebertaV3MaskedLM)

# Check that output matches.
restored_output = restored_model.predict(self.raw_batch)
self.assertAllClose(model_output, restored_output)
self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
4 changes: 3 additions & 1 deletion keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,9 @@ def test_saved_model(self, save_format, filename):
outputs = self.preprocessor(inputs)
model = keras.Model(inputs, outputs)
path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)
restored_model = keras.models.load_model(path)
self.assertAllEqual(
model(input_data)["token_ids"],
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,9 @@ def test_saved_model(self, save_format, filename):
model = keras.Model(inputs, outputs)

path = os.path.join(self.get_temp_dir(), filename)
model.save(path, save_format=save_format)
# Don't save traces in the tf format, we check compilation elsewhere.
kwargs = {"save_traces": False} if save_format == "tf" else {}
model.save(path, save_format=save_format, **kwargs)

restored_model = keras.models.load_model(path)
self.assertAllEqual(
Expand Down
Loading