Unverified Commit d65af7d8 authored by Hongkun Yu's avatar Hongkun Yu Committed by GitHub
Browse files

Merged commit includes the following changes: (#7323)

260580119  by hongkuny<hongkuny@google.com>:

    Adds expect_partial()

--

PiperOrigin-RevId: 260580119
parent 803f833c
...@@ -66,7 +66,7 @@ def export_pretraining_checkpoint( ...@@ -66,7 +66,7 @@ def export_pretraining_checkpoint(
"""Exports BERT model for as a checkpoint without optimizer. """Exports BERT model for as a checkpoint without optimizer.
Arguments: Arguments:
checkpoint_dir: Path to where training mdoel checkpoints are stored. checkpoint_dir: Path to where training model checkpoints are stored.
model: Keras model object to export. model: Keras model object to export.
checkpoint_name: File name or suffix path to export pretrained checkpoint. checkpoint_name: File name or suffix path to export pretrained checkpoint.
...@@ -83,7 +83,8 @@ def export_pretraining_checkpoint( ...@@ -83,7 +83,8 @@ def export_pretraining_checkpoint(
assert latest_checkpoint_file assert latest_checkpoint_file
logging.info('Checkpoint file %s found and restoring from ' logging.info('Checkpoint file %s found and restoring from '
'checkpoint', latest_checkpoint_file) 'checkpoint', latest_checkpoint_file)
checkpoint.restore(latest_checkpoint_file).assert_existing_objects_matched() status = checkpoint.restore(latest_checkpoint_file)
status.assert_existing_objects_matched().expect_partial()
saved_path = checkpoint.save(os.path.join(checkpoint_dir, checkpoint_name)) saved_path = checkpoint.save(os.path.join(checkpoint_dir, checkpoint_name))
logging.info('Exporting the model as a new TF checkpoint: %s', saved_path) logging.info('Exporting the model as a new TF checkpoint: %s', saved_path)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment