def save_model(output_dir: str, estimator: tf.estimator.Estimator) -> None:
def serving_input_receiver_fn() -> Any:
embeddings = tf.placeholder(dtype=tf.float32,
shape=[None, None, 768],
name='embeddings')
input_mask = tf.placeholder(dtype=tf.int32,
shape=[None, None],
name='input_mask')
features = {'embeddings': embeddings, 'input_mask': input_mask}
return tf.estimator.export.ServingInputReceiver(
features=features, receiver_tensors=features)
estimator.export_saved_model(output_dir, serving_input_receiver_fn)
def train_estimator(estimator: tf.estimator.Estimator, input_config,
train_config, export_config, task_config: TrainTaskConfig):
example_config = input_config['example_config']
label_col = input_config['label_col']
feature_specs = parse_feature_specs(example_config)
dataset_fn = get_dataset_fn(feature_specs=feature_specs,
label_col=label_col,
**train_config)
train_spec = tf.estimator.TrainSpec(dataset_fn)
eval_spec = tf.estimator.EvalSpec(dataset_fn, steps=1)
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
feature_placeholders = get_feature_placeholders(**export_config)
serving_input_receiver_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(
feature_placeholders)
if (task_config.task_type == 'chief' and task_config.task_index == 0) or \
(task_config.num_workers == 1):
logging.info("Start exporting...")
estimator.export_saved_model(
task_config.saved_model_dir,
serving_input_receiver_fn=serving_input_receiver_fn)
logging.info("Finish exporting.")
def export(self, estimator: tf.estimator.Estimator):
features = {field.name: field.as_placeholder(batch=True) for field in self.fields}
return estimator.export_saved_model(
self.path_saved_model, tf.estimator.export.build_raw_serving_input_receiver_fn(features)
)