View source on GitHub |
Pipeline for multi-task training.
Inherits From: ModelFitPipeline
, AbstractPipeline
tfr.keras.pipeline.MultiTaskPipeline(
model_builder: tfr.keras.model.AbstractModelBuilder
,
dataset_builder: tfr.keras.pipeline.AbstractDatasetBuilder
,
hparams: tfr.keras.pipeline.PipelineHparams
)
This handles a set of losses and labels. It is intended to mainly work with
MultiLabelDatasetBuilder
.
Use subclassing to customize the losses and metrics.
Example usage:
context_feature_spec = {}
example_feature_spec = {
"example_feature_1": tf.io.FixedLenFeature(
shape=(1,), dtype=tf.float32, default_value=0.0)
}
mask_feature_name = "list_mask"
label_spec_tuple = ("utility",
tf.io.FixedLenFeature(
shape=(1,),
dtype=tf.float32,
default_value=_PADDING_LABEL))
label_spec = {"task1": label_spec_tuple, "task2": label_spec_tuple}
weight_spec = ("weight",
tf.io.FixedLenFeature(
shape=(1,), dtype=tf.float32, default_value=1.))
dataset_hparams = DatasetHparams(
train_input_pattern="train.dat",
valid_input_pattern="valid.dat",
train_batch_size=128,
valid_batch_size=128)
pipeline_hparams = PipelineHparams(
model_dir="model/",
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)
model_builder = MultiTaskModelBuilder(...)
dataset_builder = MultiLabelDatasetBuilder(
context_feature_spec,
example_feature_spec,
mask_feature_name,
label_spec,
dataset_hparams,
sample_weight_spec=weight_spec)
pipeline = MultiTaskPipeline(model_builder, dataset_builder, pipeline_hparams)
pipeline.train_and_validate(verbose=1)
Methods
build_callbacks
build_callbacks() -> List[tf.keras.callbacks.Callback]
Sets up Callbacks.
Example usage:
model_builder = ModelBuilder(...)
dataset_builder = DatasetBuilder(...)
hparams = PipelineHparams(...)
pipeline = BasicModelFitPipeline(model_builder, dataset_builder, hparams)
callbacks = pipeline.build_callbacks()
Returns | |
---|---|
A list of tf.keras.callbacks.Callback or a
tf.keras.callbacks.CallbackList for tensorboard and checkpoint.
|
build_loss
build_loss() -> Dict[str, tf.keras.losses.Loss]
See AbstractPipeline
.
build_metrics
build_metrics() -> Dict[str, List[tf.keras.metrics.Metric]]
See AbstractPipeline
.
build_weighted_metrics
build_weighted_metrics() -> Dict[str, List[tf.keras.metrics.Metric]]
See AbstractPipeline
.
export_saved_model
export_saved_model(
model: tf.keras.Model,
export_to: str,
checkpoint: Optional[tf.train.Checkpoint] = None
)
Exports the trained model with signatures.
Example usage:
model_builder = ModelBuilder(...)
dataset_builder = DatasetBuilder(...)
hparams = PipelineHparams(...)
pipeline = BasicModelFitPipeline(model_builder, dataset_builder, hparams)
pipeline.export_saved_model(model_builder.build(), 'saved_model/')
Args | |
---|---|
model
|
Model to be saved. |
export_to
|
Specifies the directory the model is be exported to. |
checkpoint
|
If given, export the model with weights from this checkpoint. |
train_and_validate
train_and_validate(
verbose=0
)
Main function to train the model with TPU strategy.
Example usage:
context_feature_spec = {}
example_feature_spec = {
"example_feature_1": tf.io.FixedLenFeature(
shape=(1,), dtype=tf.float32, default_value=0.0)
}
mask_feature_name = "list_mask"
label_spec = {
"utility": tf.io.FixedLenFeature(
shape=(1,), dtype=tf.float32, default_value=0.0)
}
dataset_hparams = DatasetHparams(
train_input_pattern="train.dat",
valid_input_pattern="valid.dat",
train_batch_size=128,
valid_batch_size=128)
pipeline_hparams = pipeline.PipelineHparams(
model_dir="model/",
num_epochs=2,
steps_per_epoch=5,
validation_steps=2,
learning_rate=0.01,
loss="softmax_loss")
model_builder = SimpleModelBuilder(
context_feature_spec, example_feature_spec, mask_feature_name)
dataset_builder = SimpleDatasetBuilder(
context_feature_spec,
example_feature_spec,
mask_feature_name,
label_spec,
dataset_hparams)
pipeline = BasicModelFitPipeline(
model_builder, dataset_builder, pipeline_hparams)
pipeline.train_and_validate(verbose=1)
Args | |
---|---|
verbose
|
An int for the verbosity level. |