Unverified Commit 4c35c8d8 authored by Matt's avatar Matt Committed by GitHub
Browse files

Experimenting with adding proper get_config() and from_config() methods (#14361)

* Experimenting with adding proper get_config() and from_config() methods

* Adding a test for get/from config

* Fix test for get/from config
parent b1dbdf22
......@@ -692,6 +692,13 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin, Pu
self.config = config
self.name_or_path = config.name_or_path
def get_config(self):
return self.config
@classmethod
def from_config(cls, config, **kwargs):
return cls._from_config(config, **kwargs)
@classmethod
def _from_config(cls, config, **kwargs):
"""
......
......@@ -163,6 +163,20 @@ class TFModelTesterMixin:
self.assert_outputs_same(after_outputs, outputs)
def test_save_load_config(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(self._prepare_for_class(inputs_dict, model_class))
new_model = model_class.from_config(model.get_config())
_ = new_model(self._prepare_for_class(inputs_dict, model_class)) # Build model
new_model.set_weights(model.get_weights())
after_outputs = new_model(self._prepare_for_class(inputs_dict, model_class))
self.assert_outputs_same(after_outputs, outputs)
@tooslow
def test_graph_mode(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment