def test_sampling_model_trainer(self):
tf.random.set_seed(1)
embeddings_config = EmbeddingsConfig(
entities_count=3, relations_count=2, embeddings_dimension=4,
pretrained_entity_embeddings=self.init_entity_embeddings,
pretrained_relation_embeddings=self.init_relation_embeddings,
)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
loss_object = NormLossObject(order=2, margin=1.0)
learning_rate_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=1e-3, decay_steps=1, decay_rate=0.5
)
transe_model = TranseModel(embeddings_config, model_config)
model_trainer = SamplingModelTrainer(transe_model, loss_object, learning_rate_schedule)
dataset = SamplingEdgeDataset(
dataset_id="dataset1", inference_mode=False, dataset_type=DatasetType.TRAINING,
data_directory=self.DATASET_PATH, batch_size=2
)
samples_iterator = iter(dataset.samples)
model_trainer.train_step(training_samples=next(samples_iterator), training_step=1)
model_trainer.train_step(training_samples=next(samples_iterator), training_step=2)
self.assertEqual(2, model_trainer.optimizer.iterations)
expected_kernel = np.array([[[[1.0]], [[1.0]], [[-1.0]]]])
self.assertAllEqual(expected_kernel, transe_model.conv_layers[0].kernel.numpy())
embeddings_layer = transe_model.embeddings_layer
self.assertGreater(np.sum(self.init_entity_embeddings != embeddings_layer.entity_embeddings), 0)
self.assertGreater(np.sum(self.init_relation_embeddings != embeddings_layer.relation_embeddings), 0)
def test_non_trainable_convolutions(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
transe_model = TranseModel(embeddings_config, model_config)
transe_model.embeddings_layer.entity_embeddings.assign(
self.entity_embeddings)
transe_model.embeddings_layer.relation_embeddings.assign(
self.relation_embeddings)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
with tf.GradientTape() as gradient_tape:
model_outputs = transe_model(self.model_inputs, training=True)
loss = tf.keras.losses.mean_squared_error(
tf.zeros_like(model_outputs), model_outputs)
gradients = gradient_tape.gradient(loss,
transe_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, transe_model.trainable_variables))
expected_kernel = np.array([[[[1.0]], [[1.0]], [[-1.0]]]])
self.assertAllEqual(expected_kernel,
transe_model.conv_layers[0].kernel.numpy())
expected_bias = np.array([0.0])
self.assertAllEqual(expected_bias,
transe_model.conv_layers[0].bias.numpy())
def test_sampling_model_evaluator(self, file_writer_mock,
summary_scalar_histogram,
summary_scalar_patch):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
transe_model = TranseModel(embeddings_config, model_config)
loss_object = NormLossObject(order=2, margin=1.0)
learning_rate_scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=1e-3, decay_steps=1, decay_rate=0.5)
dataset = SamplingEdgeDataset(
dataset_id="dataset1",
inference_mode=True,
dataset_type=DatasetType.TRAINING,
data_directory=self.DATASET_PATH,
batch_size=5,
)
model_evaluator = SamplingModelEvaluator(
model=transe_model,
loss_object=loss_object,
dataset=dataset,
output_directory="logs",
learning_rate_scheduler=learning_rate_scheduler)
model_evaluator.evaluation_step(step=0)
model_evaluator.evaluation_step(step=1)
model_evaluator.log_metrics(logger=mock.MagicMock())
file_writer_mock.assert_called_once()
self.assertEqual(4, summary_scalar_histogram.call_count)
self.assertEqual(24, summary_scalar_patch.call_count)
def test_pretrained_embeddings(self):
embeddings_config = EmbeddingsConfig(
entities_count=3,
relations_count=2,
pretrained_entity_embeddings=self.entity_embeddings,
embeddings_dimension=4,
pretrained_relation_embeddings=self.relation_embeddings)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
transe_model = TranseModel(embeddings_config, model_config)
outputs = transe_model(self.model_inputs, training=True)
self.assertEqual((2, 4), outputs.shape)
self.assertAllClose([[2., 2., 1., 2.], [3., 2., 4., 3.]],
outputs.numpy())
def test_dropout_layer(self):
tf.random.set_seed(1)
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=False,
dropout_rate=0.999)
transe_model = TranseModel(embeddings_config, model_config)
transe_model.embeddings_layer.entity_embeddings.assign(
self.entity_embeddings)
transe_model.embeddings_layer.relation_embeddings.assign(
self.relation_embeddings)
outputs = transe_model(self.model_inputs, training=True)
self.assertEqual((2, 4), outputs.shape)
self.assertAllClose(tf.zeros_like(outputs), outputs)
def test_normalize_embeddings(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=2)
model_config = ConvModelConfig(include_reduce_dim_layer=False,
normalize_embeddings=True)
transe_model = TranseModel(embeddings_config, model_config)
transe_model.embeddings_layer.entity_embeddings.assign([[0., 0.],
[1., 0.],
[0.8, 0.6]])
transe_model.embeddings_layer.relation_embeddings.assign([[1.6, 1.2],
[2.4, 3.2]])
outputs = transe_model(self.model_inputs, training=True)
self.assertEqual((2, 2), outputs.shape)
self.assertAllClose([[-0.2, 0.6], [0.8, 0.2]], outputs.numpy())
def _create_model(embeddings_config: EmbeddingsConfig, model_type: ModelType):
type_mappings = {
ModelType.TRANSE:
lambda: TranseModel(embeddings_config),
ModelType.STRANSE:
lambda: STranseModel(embeddings_config),
ModelType.CONVKB:
lambda: ConvKBModel(embeddings_config),
ModelType.CONVE:
lambda: ConvEModel(embeddings_config),
ModelType.TRANSFORMER_SOFTMAX:
lambda: TransformerSoftmaxModel(embeddings_config),
}
if model_type not in type_mappings:
raise ValueError(f"Invalid model type: {model_type}")
return type_mappings[model_type]()
def test_reduce_dim_layer(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=True)
transe_model = TranseModel(embeddings_config, model_config)
transe_model.embeddings_layer.entity_embeddings.assign(
self.entity_embeddings)
transe_model.embeddings_layer.relation_embeddings.assign(
self.relation_embeddings)
transe_model(self.model_inputs, training=True)
transe_model.reduce_layer.set_weights(
[np.array([[1.], [2.], [1.], [0.]]),
np.array([-10.])])
output = transe_model(self.model_inputs, training=True)
self.assertEqual((2, 1), output.shape)
self.assertAllClose([[-3.], [1.]], output.numpy())
def test_gin_config(self):
gin_config = """
ConvModelConfig.include_reduce_dim_layer = False
ConvModelConfig.normalize_embeddings = False
ConvModelConfig.dropout_rate = 0.0
TranseModel.model_config = @ConvModelConfig()
"""
gin.parse_config(gin_config)
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
transe_model = TranseModel(embeddings_config)
transe_model.embeddings_layer.entity_embeddings.assign(
self.entity_embeddings)
transe_model.embeddings_layer.relation_embeddings.assign(
self.relation_embeddings)
outputs = transe_model(self.model_inputs, training=True)
self.assertEqual((2, 4), outputs.shape)
self.assertAllClose([[2., 2., 1., 2.], [3., 2., 4., 3.]],
outputs.numpy())
def test_trainable_embeddings(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
transe_model = TranseModel(embeddings_config, model_config)
transe_model.embeddings_layer.entity_embeddings.assign(
self.entity_embeddings)
transe_model.embeddings_layer.relation_embeddings.assign(
self.relation_embeddings)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
with tf.GradientTape() as gradient_tape:
model_outputs = transe_model(self.model_inputs, training=True)
loss = tf.keras.losses.mean_squared_error(
tf.zeros_like(model_outputs), model_outputs)
gradients = gradient_tape.gradient(loss,
transe_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, transe_model.trainable_variables))
self.assertGreater(
np.sum(self.entity_embeddings !=
transe_model.embeddings_layer.entity_embeddings.numpy()), 0)