No public description

PiperOrigin-RevId: 712914895

Author: tensorflower-gardener

Diff for: official/projects/pix2seq/configs/pix2seq.py

@@ -115,6 +115,11 @@ class BackboneConfig(hyperparams.Config):
   )
   # Optional checkpoint to load for this backbone.
   init_checkpoint: Optional[str] = None
+  # If loading an init_checkpoint, whether to assert that all objects in the
+  # Python program are matched by the checkpoint.
+  # If False, understand that only the weak assertion of a non-trivial match
+  # will be made.
+  assert_existing_objects_matched: bool = True
 
 
 @dataclasses.dataclass

Diff for: official/projects/pix2seq/tasks/pix2seq_task.py

@@ -141,8 +141,13 @@ def initialize(self, model: tf_keras.Model):
           backbone.load_checkpoint(ckpt_filepath=backbone_init_ckpt)
         else:
           ckpt = tf.train.Checkpoint(backbone=backbone)
-          status = ckpt.restore(backbone_init_ckpt)
-          status.expect_partial().assert_existing_objects_matched()
+          status = (
+              ckpt.restore(backbone_init_ckpt)
+              .expect_partial()
+              .assert_nontrivial_match()
+          )
+          if backbone_config.assert_existing_objects_matched:
+            status.assert_existing_objects_matched()
 
         logging.info(
             'Finished loading pretrained backbone from %s', backbone_init_ckpt