def test_gam_model_builder_with_spec(self):
gam_scorer = model_lib.GAMScorer(
example_hidden_layer_dims=[10, 10], context_hidden_layer_dims=[10, 10])
gam_model = model_lib.ModelBuilder(
input_creator=self._input_creator,
preprocessor=self._preprocessor,
mask_feature_name="mask",
scorer=gam_scorer,
name="test_model").build()
output = gam_model({
"context_1": tf.convert_to_tensor([[1]]),
"feature_1": tf.convert_to_tensor([[[0.], [1], [2]]]),
"feature_2": tf.convert_to_tensor([[[0.], [1], [2]]]),
"mask": tf.convert_to_tensor([[True, True, True]]),
})
self.assertAllEqual(output.shape.as_list(), [1, 3])
def test_dnn_model_builder_with_spec(self):
dnn_scorer = model_lib.DNNScorer(hidden_layer_dims=[10, 10], output_units=1)
dnn_model = model_lib.ModelBuilder(
input_creator=self._input_creator,
preprocessor=self._preprocessor,
scorer=dnn_scorer,
mask_feature_name="mask",
name="test_model",
).build()
output = dnn_model({
"context_1": tf.ragged.constant([[3, 1, 4, 1]]),
"feature_1": tf.convert_to_tensor([[[0.], [1], [2]]]),
"feature_2": tf.convert_to_tensor([[[0.], [1], [2]]]),
"mask": tf.convert_to_tensor([[True, True, True]]),
})
self.assertAllEqual(output.shape.as_list(), [1, 3])
def test_model_to_saved_model_dense_inputs(self):
dnn_scorer = model_lib.DNNScorer(hidden_layer_dims=[10, 10],
output_units=1)
dnn_model = model_lib.ModelBuilder(
input_creator=self._input_creator,
preprocessor=self._preprocessor,
scorer=dnn_scorer,
mask_feature_name="mask",
name="test_model",
).build()
@tf.function(
input_signature=[tf.TensorSpec(shape=(1, ), dtype=tf.string)])
def _predict(serialized):
features = data.parse_from_example_list(
serialized,
context_feature_spec=self._context_feature_spec,
example_feature_spec=self._example_feature_spec,
mask_feature_name="mask")
scores = dnn_model(inputs=features, training=False)
return {"predictions": scores}
dnn_model.infer_from_proto = _predict
# Export the model to a SavedModel.
tf.saved_model.save(dnn_model,
export_dir="/tmp/keras_model",
signatures={"predict": dnn_model.infer_from_proto})
# Import ranker from SavedModel.
imported = tf.saved_model.load("/tmp/keras_model")
imported_model = imported.signatures["predict"]
imported_output = imported_model(
tf.convert_to_tensor([ELWC_PROTO.SerializeToString()
]))["predictions"]
output = dnn_model({
"context_1":
tf.ragged.constant([[3, 1, 4, 1]]),
"feature_1":
tf.convert_to_tensor([[[0.], [1], [2]]]),
"feature_2":
tf.convert_to_tensor([[[0.], [1], [2]]]),
"mask":
tf.convert_to_tensor([[True, True, True]]),
})
self.assertAllClose(output.numpy(), imported_output.numpy())
def test_multi_task_scorer(self):
mt_model = model_lib.ModelBuilder(
input_creator=self._input_creator,
preprocessor=self._preprocessor,
scorer=DummyMultiTaskScorer(),
mask_feature_name="mask",
name="test_model",
).build()
output = mt_model({
"context_1": tf.ragged.constant([[3, 1, 4, 1]]),
"feature_1": tf.convert_to_tensor([[[0.], [1], [2]]]),
"feature_2": tf.convert_to_tensor([[[0.], [1], [2]]]),
"mask": tf.convert_to_tensor([[True, True, True]]),
})
self.assertEqual(len(output), 2)
self.assertAllEqual(output["task1"].numpy(), [[1., 2., 3.]])
self.assertAllEqual(output["task2"].numpy(), [[2., 4., 6.]])
def test_pipeline_with_multi_task(self):
data_dir = self.create_tempdir()
data_file = os.path.join(data_dir, "elwc.tfrecord")
if tf.io.gfile.exists(data_file):
tf.io.gfile.remove(data_file)
with tf.io.TFRecordWriter(data_file) as writer:
for _ in range(256):
writer.write(ELWC.SerializeToString())
model_dir = os.path.join(data_dir, "model")
dataset_hparams = pipeline.DatasetHparams(
train_input_pattern=data_file,
valid_input_pattern=data_file,
train_batch_size=128,
valid_batch_size=128,
list_size=2,
dataset_reader=tf.data.TFRecordDataset,
convert_labels_to_binary=False)
pipeline_hparams = pipeline.PipelineHparams(
model_dir=model_dir,
num_epochs=2,
steps_per_epoch=5,
validation_steps=2,
learning_rate=0.01,
loss={
"task1": "softmax_loss",
"task2": "pairwise_logistic_loss"
},
loss_weights={
"task1": 1.0,
"task2": 2.0
},
export_best_model=True,
strategy="MirroredStrategy")
context_feature_spec = {
"cf_1":
tf.io.FixedLenFeature(
shape=(1,), dtype=tf.float32, default_value=0.0),
}
example_feature_spec = {
"custom_features_{}".format(i + 1):
tf.io.FixedLenFeature(shape=(1,), dtype=tf.float32, default_value=0.0)
for i in range(3)
}
label_spec = (_LABEL_FEATURE,
tf.io.FixedLenFeature(
shape=(1,),
dtype=tf.float32,
default_value=_PADDING_LABEL))
label_spec = {"task1": label_spec, "task2": label_spec}
weight_spec = ("weight",
tf.io.FixedLenFeature(
shape=(1,), dtype=tf.float32, default_value=1.))
model_builder = model_lib.ModelBuilder(
input_creator=model_lib.FeatureSpecInputCreator(context_feature_spec,
example_feature_spec),
preprocessor=model_lib.PreprocessorWithSpec({}),
scorer=DummyMultiTaskScorer(),
mask_feature_name=_MASK,
name="multi_task_model",
)
ranking_pipeline = pipeline.MultiTaskPipeline(
model_builder,
dataset_builder=pipeline.MultiLabelDatasetBuilder(
context_feature_spec,
example_feature_spec,
_MASK,
label_spec,
dataset_hparams,
sample_weight_spec=weight_spec),
hparams=pipeline_hparams)
ranking_pipeline.train_and_validate(verbose=1)
latest_model_path = os.path.join(model_dir, "export/latest_model")
self.assertTrue(tf.saved_model.contains_saved_model(latest_model_path))
latest_model = tf.saved_model.load(export_dir=latest_model_path)
listwise_predictor = latest_model.signatures[
tf.saved_model.PREDICT_METHOD_NAME]
listwise_logits = listwise_predictor(
tf.convert_to_tensor([ELWC.SerializeToString()] * 2))["task1"]
self.assertAllEqual([2, 2], listwise_logits.get_shape().as_list())
pointwise_predictor = latest_model.signatures[
tf.saved_model.REGRESS_METHOD_NAME]
pointwise_logits = pointwise_predictor(
tf.convert_to_tensor([
EXAMPLE_PROTO_1.SerializeToString(),
EXAMPLE_PROTO_2.SerializeToString()
]))["task1"]
self.assertAllEqual([2], pointwise_logits.get_shape().as_list())
self.assertAllClose(pointwise_logits, listwise_logits[0])
def test_pipeline_with_feature_specs(self, convert_labels_to_binary):
data_dir = self.create_tempdir()
data_file = os.path.join(data_dir, "elwc.tfrecord")
if tf.io.gfile.exists(data_file):
tf.io.gfile.remove(data_file)
with tf.io.TFRecordWriter(data_file) as writer:
for _ in range(256):
writer.write(ELWC.SerializeToString())
context_feature_spec = {
"cf_1":
tf.io.FixedLenFeature(shape=(1, ),
dtype=tf.float32,
default_value=0.0),
}
example_feature_spec = {
"custom_features_{}".format(i + 1):
tf.io.FixedLenFeature(shape=(1, ),
dtype=tf.float32,
default_value=0.0)
for i in range(3)
}
label_spec = (_LABEL_FEATURE,
tf.io.FixedLenFeature(shape=(1, ),
dtype=tf.float32,
default_value=_PADDING_LABEL))
train_data_config = task_lib.RankingDataConfig(
input_path=data_file,
is_training=True,
global_batch_size=128,
list_size=2,
mask_feature_name=_MASK,
dataset_fn="tfrecord",
convert_labels_to_binary=convert_labels_to_binary)
validation_data_config = task_lib.RankingDataConfig(
input_path=data_file,
is_training=False,
global_batch_size=128,
list_size=2,
mask_feature_name=_MASK,
dataset_fn="tfrecord",
convert_labels_to_binary=convert_labels_to_binary)
dnn_scorer = model_lib.DNNScorer(hidden_layer_dims=[16, 8],
output_units=1)
model_builder = model_lib.ModelBuilder(
input_creator=model_lib.FeatureSpecInputCreator(
context_feature_spec, example_feature_spec),
preprocessor=model_lib.PreprocessorWithSpec({}),
scorer=dnn_scorer,
mask_feature_name=_MASK,
name="test_model",
)
ranking_task_config = task_lib.RankingTaskConfig(
train_data=train_data_config,
validation_data=validation_data_config,
loss="softmax_loss")
task = task_lib.RankingTask(ranking_task_config,
model_builder=model_builder,
context_feature_spec=context_feature_spec,
example_feature_spec=example_feature_spec,
label_spec=label_spec)
model = task.build_model()
metrics = task.build_metrics()
train_dataset = task.build_inputs(ranking_task_config.train_data)
vali_dataset = task.build_inputs(ranking_task_config.validation_data)
task.initialize(model)
train_iterator = iter(train_dataset)
vali_iterator = iter(vali_dataset)
optimizer = tf.keras.optimizers.SGD(lr=0.1)
task.train_step(next(train_iterator),
model,
optimizer,
metrics=metrics)
task.validation_step(next(vali_iterator), model, metrics=metrics)
