From bb04edb45b73d06bc8674408df193d94c3bf77a1 Mon Sep 17 00:00:00 2001 From: thomwolf Date: Thu, 10 Oct 2019 13:08:24 +0200 Subject: [PATCH] Add tests that TF 2.0 model can be integrated with other Keras modules --- transformers/tests/modeling_tf_common_test.py | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/transformers/tests/modeling_tf_common_test.py b/transformers/tests/modeling_tf_common_test.py index 49a5776e69..5b21ad15d7 100644 --- a/transformers/tests/modeling_tf_common_test.py +++ b/transformers/tests/modeling_tf_common_test.py @@ -22,6 +22,7 @@ import random import shutil import unittest import uuid +import tempfile import pytest import sys @@ -36,6 +37,20 @@ if is_tf_available(): else: pytestmark = pytest.mark.skip("Require TensorFlow") +if sys.version_info[0] == 2: + import cPickle as pickle + + class TemporaryDirectory(object): + """Context manager for tempfile.mkdtemp() so it's usable with "with" statement.""" + def __enter__(self): + self.name = tempfile.mkdtemp() + return self.name + def __exit__(self, exc_type, exc_value, traceback): + shutil.rmtree(self.name) +else: + import pickle + TemporaryDirectory = tempfile.TemporaryDirectory + unicode = str def _config_zero_init(config): configs_no_init = copy.deepcopy(config) @@ -66,13 +81,25 @@ class TFCommonTestCases: # self.assertIn(param.data.mean().item(), [0.0, 1.0], # msg="Parameter {} of model {} seems not properly initialized".format(name, model_class)) + def test_save_load(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + outputs = model(inputs_dict) + + with TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname) + model = model_class.from_pretrained(tmpdirname) + after_outputs = model(inputs_dict) + max_diff = np.amax(np.abs(after_outputs[0].numpy() - outputs[0].numpy())) + self.assertLessEqual(max_diff, 1e-5) def test_pt_tf_model_equivalence(self): if not is_torch_available(): return import torch - import numpy as np import transformers config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -99,6 +126,34 @@ class TFCommonTestCases: max_diff = np.amax(np.abs(tfo[0].numpy() - pto[0].numpy())) self.assertLessEqual(max_diff, 2e-2) + def test_compile_tf_model(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + input_ids = tf.keras.Input(batch_shape=(2, 2000), name='input_ids', dtype='int32') + optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0) + loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) + metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy') + + for model_class in self.all_model_classes: + # Prepare our model + model = model_class(config) + + # Let's load it from the disk to be sure we can use pretrained weights + with TemporaryDirectory() as tmpdirname: + outputs = model(inputs_dict) # build the model + model.save_pretrained(tmpdirname) + model = model_class.from_pretrained(tmpdirname) + + outputs_dict = model(input_ids) + hidden_states = outputs_dict[0] + + # Add a dense layer on top to test intetgration with other keras modules + outputs = tf.keras.layers.Dense(2, activation='softmax', name='outputs')(hidden_states) + + # Compile extended model + extended_model = tf.keras.Model(inputs=[input_ids], outputs=[outputs]) + extended_model.compile(optimizer=optimizer, loss=loss, metrics=[metric]) + def test_keyword_and_dict_args(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()