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_filter_heights(self):
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convolutions_config = ConvolutionsConfig(
filter_heights=[1, 2],
filters_count_per_height=1,
activation="relu",
use_constant_initialization=False)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config, convolutions_config)
convkb_model.embeddings_layer.entity_embeddings.assign(
self.default_entity_embeddings)
convkb_model.embeddings_layer.relation_embeddings.assign(
self.default_relation_embeddings)
convkb_model(self.model_inputs)
convkb_model.conv_layers[0].kernel.assign(self.default_kernel)
convkb_model.conv_layers[0].bias.assign(self.default_bias)
convkb_model.conv_layers[1].kernel.assign([[[[1.0]], [[1.0]],
[[-1.0]]],
[[[-1.0]], [[0.0]],
[[2.0]]]])
convkb_model.conv_layers[1].bias.assign([1.0])
outputs = convkb_model(self.model_inputs, training=True)
self.assertEqual((2, 7), outputs.shape)
self.assertAllClose(
[[2., 2., 1., 2., 5., 7., 4.], [3., 2., 4., 3., 7., 5., 8.]],
outputs.numpy())
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_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_outputs(self):
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config,
self.default_convolutions_config)
self._set_default_weights(convkb_model)
outputs = convkb_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 setUp(self):
self.entity_embeddings = np.array([[0., 0., 0., 0.], [0.1, 0.2, 0.3, 0.4], [3., 4., 0., 0.]], dtype=np.float32)
self.relation_embeddings = np.array([[3., 3., 3., 3.], [0.6, -0.8, 0.0, 0.0]], dtype=np.float32)
self.embeddings_config = EmbeddingsConfig(entities_count=3, relations_count=2, embeddings_dimension=4)
self.model_config = ConvModelConfig(include_reduce_dim_layer=False)
edge_object_type = [ObjectType.ENTITY.value, ObjectType.RELATION.value, ObjectType.ENTITY.value]
self.model_inputs = {
"object_ids": tf.constant([[0, 0, 1], [0, 1, 2]], dtype=tf.int32),
"object_types": tf.constant([edge_object_type, edge_object_type], dtype=tf.int32),
}
gin.clear_config()
def test_dropout_layer(self):
tf.random.set_seed(1)
model_config = ConvModelConfig(include_reduce_dim_layer=False,
dropout_rate=0.999)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config,
self.default_convolutions_config)
self._set_default_weights(convkb_model)
outputs = convkb_model(self.model_inputs, training=True)
self.assertEqual((2, 4), outputs.shape)
self.assertAllClose(tf.zeros_like(outputs), outputs)
def test_reduce_dim_layer(self):
model_config = ConvModelConfig(include_reduce_dim_layer=True)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config,
self.default_convolutions_config)
self._set_default_weights(convkb_model)
convkb_model.reduce_layer.set_weights(
[np.array([[1.], [2.], [1.], [0.]]),
np.array([-10.])])
outputs = convkb_model(self.model_inputs, training=True)
self.assertEqual((2, 1), outputs.shape)
self.assertAllClose([[-3.], [1.]], outputs.numpy())
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_non_constant_initialization(self):
tf.random.set_seed(1)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convolutions_config = ConvolutionsConfig(
filter_heights=[1],
filters_count_per_height=1,
activation="relu",
use_constant_initialization=False)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config, convolutions_config)
convkb_model(self.model_inputs)
unequal_elements = convkb_model.conv_layers[0].kernel.numpy(
) != np.array([[[[0.1]], [[0.1]], [[-0.1]]]])
self.assertGreater(np.sum(unequal_elements), 0)
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 test_pretrained_embeddings(self):
embeddings_config = EmbeddingsConfig(
entities_count=3,
relations_count=2,
embeddings_dimension=4,
pretrained_entity_embeddings=self.default_entity_embeddings,
pretrained_relation_embeddings=self.default_relation_embeddings)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convkb_model = ConvKBModel(embeddings_config, model_config,
self.default_convolutions_config)
self._set_default_weights(convkb_model)
outputs = convkb_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_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_activation(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convolutions_config = ConvolutionsConfig(
filter_heights=[1],
filters_count_per_height=1,
activation="relu",
use_constant_initialization=False)
convkb_model = ConvKBModel(embeddings_config, model_config,
convolutions_config)
self._set_default_weights(convkb_model)
convkb_model.embeddings_layer.relation_embeddings.assign(
-self.default_relation_embeddings)
outputs = convkb_model(self.model_inputs, training=True)
self.assertEqual((2, 4), outputs.shape)
self.assertAllClose(tf.zeros_like(outputs), outputs.numpy())
def test_model(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
model = TransformerTranseModel(embeddings_config, model_config)
edge_object_type = [
ObjectType.ENTITY.value, ObjectType.RELATION.value,
ObjectType.ENTITY.value, ObjectType.ENTITY.value,
ObjectType.RELATION.value
]
inputs = {
"object_ids":
tf.constant([[0, 0, 1, 1, 0], [1, 1, 2, 0, 1]], dtype=tf.int32),
"object_types":
tf.constant([edge_object_type, edge_object_type], dtype=tf.int32),
}
outputs = model(inputs)
self.assertAllEqual((2, 4), outputs.shape)
def test_constant_initialization(self):
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convolutions_config = ConvolutionsConfig(
filter_heights=[1],
filters_count_per_height=1,
activation="relu",
use_constant_initialization=True)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config, convolutions_config)
convkb_model.embeddings_layer.entity_embeddings.assign(
self.default_entity_embeddings)
convkb_model.embeddings_layer.relation_embeddings.assign(
self.default_relation_embeddings)
convkb_model(self.model_inputs)
convkb_model.conv_layers[0].bias.assign(self.default_bias)
outputs = convkb_model(self.model_inputs, training=True)
self.assertEqual((2, 4), outputs.shape)
self.assertAllClose([[0.2, 0.2, 0.1, 0.2], [0.3, 0.2, 0.4, 0.3]],
outputs.numpy())
def test_normalize_embeddings(self):
model_config = ConvModelConfig(include_reduce_dim_layer=False,
normalize_embeddings=True)
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=2)
convkb_model = ConvKBModel(embeddings_config, model_config,
self.default_convolutions_config)
convkb_model.embeddings_layer.entity_embeddings.assign([[0., 0.],
[1., 0.],
[0.8, 0.6]])
convkb_model.embeddings_layer.relation_embeddings.assign([[1.6, 1.2],
[2.4, 3.2]])
convkb_model(self.model_inputs)
convkb_model.conv_layers[0].kernel.assign(self.default_kernel)
convkb_model.conv_layers[0].bias.assign(self.default_bias)
outputs = convkb_model(self.model_inputs, training=True)
self.assertEqual((2, 2), outputs.shape)
self.assertAllClose([[0.0, 0.6], [0.8, 0.2]], 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)
def test_trainable_convolutions(self):
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convkb_model = ConvKBModel(self.default_embeddings_config,
model_config,
self.default_convolutions_config)
self._set_default_weights(convkb_model)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
with tf.GradientTape() as gradient_tape:
model_outputs = convkb_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,
convkb_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, convkb_model.trainable_variables))
self.assertGreater(
np.sum(self.default_kernel !=
convkb_model.conv_layers[0].kernel.numpy()), 0)
self.assertGreater(
np.sum(
self.default_bias != convkb_model.conv_layers[0].bias.numpy()),
0)
def test_non_trainable_embeddings(self):
embeddings_config = EmbeddingsConfig(entities_count=3,
relations_count=2,
embeddings_dimension=4,
trainable_embeddings=False)
model_config = ConvModelConfig(include_reduce_dim_layer=False)
convkb_model = ConvKBModel(embeddings_config, model_config,
self.default_convolutions_config)
self._set_default_weights(convkb_model)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
with tf.GradientTape() as gradient_tape:
model_outputs = convkb_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,
convkb_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, convkb_model.trainable_variables))
self.assertAllEqual(
self.default_entity_embeddings,
convkb_model.embeddings_layer.entity_embeddings.numpy())
self.assertAllEqual(
self.default_relation_embeddings,
convkb_model.embeddings_layer.relation_embeddings.numpy())