View source on GitHub |
Pipleine for single-task training.
Inherits From: ModelFitPipeline
, AbstractPipeline
tfr.keras.pipeline.SimplePipeline(
model_builder: tfr.keras.model.AbstractModelBuilder
,
dataset_builder: tfr.keras.pipeline.AbstractDatasetBuilder
,
hparams: tfr.keras.pipeline.PipelineHparams
)
This handles a single loss and works with SimpleDatasetBuilder
. This can
also work with MultiLabelDatasetBuilder
. In this case, the same loss, as
well as all metrics, will be applied to all labels and their predictions
uniformly.
Use subclassing to customize the loss 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 = {
"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 = SimplePipeline(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() -> tf.keras.losses.Loss
See AbstractPipeline
.
build_metrics
build_metrics() -> List[tf.keras.metrics.Metric]
See AbstractPipeline
.
build_weighted_metrics
build_weighted_metrics() -> 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. |