def testConfigForMultipleApplications(self): original_model = tf.keras.Sequential([ tf.keras.layers.Dense(1, activation="softmax"), ]) # Arbitrary losses. loss = {"k1": losses.MMDLoss(), "k2": losses.AbsoluteCorrelationLoss()} # Arbitrary values. loss_weight = {"k1": 3.4, "k2": 2.9} model = min_diff_model.MinDiffModel(original_model, loss, loss_weight) config = model.get_config() self.assertSetEqual( set(config.keys()), set(["original_model", "loss", "loss_weight", "name"])) self.assertDictEqual(config["loss"], loss) self.assertDictEqual(config["loss_weight"], loss_weight) # Test building the model from the config. model_from_config = min_diff_model.MinDiffModel.from_config(config) self.assertIsInstance(model_from_config, min_diff_model.MinDiffModel) self.assertIsInstance(model_from_config.original_model, tf.keras.Sequential) self.assertDictEqual(model_from_config._loss, loss) self.assertDictEqual(model_from_config._loss_weight, loss_weight)
def testSerializationWithTransformAndKernel(self): original_model = tf.keras.Sequential([ tf.keras.layers.Dense(1, activation="softmax"), ]) predictions_fn = lambda x: x * 5.1 # Arbitrary operation. loss_weight = 2.3 # Arbitrary value. model_name = "custom_model_name" # Arbitrary name. model = min_diff_model.MinDiffModel( original_model, losses.MMDLoss("laplacian"), # Non-default Kernel. loss_weight=loss_weight, predictions_transform=predictions_fn, name=model_name) serialized_model = tf.keras.utils.serialize_keras_object(model) deserialized_model = tf.keras.layers.deserialize(serialized_model) self.assertIsInstance(deserialized_model, min_diff_model.MinDiffModel) val = 7 # Arbitrary value. self.assertEqual( deserialized_model._predictions_transform(val), predictions_fn(val)) self.assertIsInstance(deserialized_model._loss, losses.MMDLoss) self.assertIsInstance(deserialized_model._loss.predictions_kernel, losses.LaplacianKernel) self.assertEqual(deserialized_model._loss_weight, loss_weight) self.assertEqual(deserialized_model.name, model_name)
def testMinDiffModelRaisesErrorWithBadKwarg(self): original_model = tf.keras.Sequential( [tf.keras.layers.Dense(1, activation="softmax")]) with self.assertRaisesRegex( TypeError, "problem initializing the MinDiffModel instance"): _ = min_diff_model.MinDiffModel( original_model, losses.MMDLoss(), bad_kwarg="some value")
def testMultipleApplicationsEvalOutputs(self): original_model = tf.keras.Sequential( [tf.keras.layers.Dense(1, activation="softmax")]) model = min_diff_model.MinDiffModel(original_model, { "k1": losses.MMDLoss(), "k2": losses.MMDLoss() }) model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["acc"]) # Evaluate with min_diff_data. output_metrics = model.test_step( iter(self.multi_min_diff_dataset).get_next()) self.assertSetEqual( set(output_metrics.keys()), set(["loss", "acc", "k1_min_diff_loss", "k2_min_diff_loss"])) # Evaluate without min_diff_data. output_metrics = model.test_step(iter(self.original_dataset).get_next()) self.assertSetEqual(set(output_metrics.keys()), set(["loss", "acc"]))
def testMultipleApplicationsWithSequentialModel(self): original_model = tf.keras.Sequential( [tf.keras.layers.Dense(1, activation="softmax")]) model = min_diff_model.MinDiffModel(original_model, { "k1": losses.MMDLoss(), "k2": losses.MMDLoss() }) model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["acc"]) history = model.fit(self.multi_min_diff_dataset) self.assertSetEqual( set(history.history.keys()), set(["loss", "acc", "k1_min_diff_loss", "k2_min_diff_loss"])) # Evaluate with min_diff_data. model.evaluate(self.multi_min_diff_dataset) # Evaluate and run inference without min_diff_data. model.evaluate(self.original_dataset) model.predict(self.original_dataset)
def testGetConfigErrorFromOriginalModel(self): class CustomModel(tf.keras.Model): pass # No need to add any other implementation for this test. original_model = CustomModel() model = min_diff_model.MinDiffModel(original_model, losses.MMDLoss()) with self.assertRaisesRegex( NotImplementedError, "MinDiffModel cannot create a config.*" "original_model.*not implemented get_config.*or has an error"): _ = model.get_config()
def testConfig(self): original_model = tf.keras.Sequential([ tf.keras.layers.Dense(1, activation="softmax"), ]) model = min_diff_model.MinDiffModel(original_model, losses.MMDLoss()) config = model.get_config() self.assertSetEqual( set(config.keys()), set(["original_model", "loss", "loss_weight", "name"])) # Test building the model from the config. model_from_config = min_diff_model.MinDiffModel.from_config(config) self.assertIsInstance(model_from_config, min_diff_model.MinDiffModel) self.assertIsInstance(model_from_config.original_model, tf.keras.Sequential)
def testWithFunctionalModel(self): inputs = tf.keras.Input(1) outputs = tf.keras.layers.Dense(1, activation="softmax")(inputs) original_model = tf.keras.Model(inputs=inputs, outputs=outputs) model = min_diff_model.MinDiffModel(original_model, losses.MMDLoss()) model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["acc"]) history = model.fit(self.min_diff_dataset) self.assertSetEqual( set(history.history.keys()), set(["loss", "acc", "min_diff_loss"])) # Evaluate with min_diff_data. model.evaluate(self.min_diff_dataset) # Evaluate and run inference without min_diff_data. model.evaluate(self.original_dataset) model.predict(self.original_dataset)
def testSerialization(self): original_model = tf.keras.Sequential([ tf.keras.layers.Dense(1, activation="softmax"), ]) loss_weight = 2.3 # Arbitrary value. model_name = "custom_model_name" # Arbitrary name. model = min_diff_model.MinDiffModel( original_model, losses.MMDLoss(), loss_weight=loss_weight, name=model_name) serialized_model = tf.keras.utils.serialize_keras_object(model) deserialized_model = tf.keras.layers.deserialize(serialized_model) self.assertIsInstance(deserialized_model, min_diff_model.MinDiffModel) self.assertIsNone(deserialized_model._predictions_transform) self.assertIsInstance(deserialized_model._loss, losses.MMDLoss) self.assertEqual(deserialized_model._loss_weight, loss_weight) self.assertEqual(deserialized_model.name, model_name)
def testSerializationForMultipleApplications(self): original_model = tf.keras.Sequential([ tf.keras.layers.Dense(1, activation="softmax"), ]) # Arbitrary losses. loss = {"k1": losses.MMDLoss(), "k2": losses.AbsoluteCorrelationLoss()} # Arbitrary values. loss_weight = {"k1": 3.4, "k2": 2.9} model_name = "custom_model_name" # Arbitrary name. model = min_diff_model.MinDiffModel( original_model, loss, loss_weight, name=model_name) serialized_model = tf.keras.utils.serialize_keras_object(model) deserialized_model = tf.keras.layers.deserialize(serialized_model) self.assertIsInstance(deserialized_model, min_diff_model.MinDiffModel) self.assertIsNone(deserialized_model._predictions_transform) self.assertDictEqual(deserialized_model._loss, loss) self.assertDictEqual(deserialized_model._loss_weight, loss_weight) self.assertEqual(deserialized_model.name, model_name)
def testConfigWithCustomBaseImplementation(self): class CustomModel(tf.keras.Model): def __init__(self, val, **kwargs): super(CustomModel, self).__init__(**kwargs) self.val = val def get_config(self): return {"val": self.val} class CustomMinDiffModel(min_diff_model.MinDiffModel, CustomModel): pass # No additional implementation needed. original_val = 4 # Arbitrary value passed in. original_model = CustomModel(original_val) min_diff_model_val = 5 # Different arbitrary value passed in. model = CustomMinDiffModel( original_model=original_model, loss=losses.MMDLoss(), val=min_diff_model_val) self.assertEqual(model.val, min_diff_model_val) config = model.get_config() self.assertSetEqual( set(config.keys()), set(["original_model", "loss", "loss_weight", "name", "val"])) self.assertEqual(config["val"], model.val) self.assertEqual(config["original_model"].val, original_model.val) # Test building the model from the config. model_from_config = CustomMinDiffModel.from_config(config) self.assertIsInstance(model_from_config, CustomMinDiffModel) self.assertIsInstance(model_from_config.original_model, CustomModel) self.assertEqual(model_from_config.val, min_diff_model_val) self.assertEqual(model_from_config.original_model.val, original_val)
def testSaveForContinuedTraining(self): original_model = tf.keras.Sequential([ tf.keras.layers.Dense(1, activation="softmax"), ]) model = min_diff_model.MinDiffModel(original_model, losses.MMDLoss()) model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["acc"]) _ = model.fit(self.min_diff_dataset) with tempfile.TemporaryDirectory() as tmp: path = os.path.join(tmp, "saved_model") model.save(path) loaded_model = tf.keras.models.load_model(path) self.assertIsInstance(loaded_model, min_diff_model.MinDiffModel) # Run more training on loaded_model. loaded_model.fit(self.min_diff_dataset) # Run inference on loaded_model. loaded_model.predict(self.original_dataset)