Internal change

PiperOrigin-RevId: 327325736

official/nlp/modeling/models/bert_classifier.py

@@ -97,7 +97,11 @@ class BertClassifier(tf.keras.Model):
 
   @property
   def checkpoint_items(self):
-    return dict(encoder=self._network)
+    items = dict(encoder=self._network)
+    if hasattr(self.classifier, 'checkpoint_items'):
+      for key, item in self.classifier.checkpoint_items.items():
+        items['.'.join([self.classifier.name, key])] = item
+    return items
 
   def get_config(self):
     return self._config

official/nlp/tasks/sentence_prediction.py

@@ -215,9 +215,8 @@ class SentencePredictionTask(base_task.Task):
     pretrain2finetune_mapping = {
         'encoder': model.checkpoint_items['encoder'],
     }
-    # TODO(b/160251903): Investigate why no pooler dense improves finetuning
-    # accuracies.
     if self.task_config.init_cls_pooler:
+      # This option is valid when use_encoder_pooler is false.
       pretrain2finetune_mapping[
           'next_sentence.pooler_dense'] = model.checkpoint_items[
               'sentence_prediction.pooler_dense']

official/nlp/tasks/sentence_prediction_test.py

@@ -87,34 +87,40 @@ class SentencePredictionTaskTest(tf.test.TestCase, parameterized.TestCase):
     task.train_step(next(iterator), model, optimizer, metrics=metrics)
     task.validation_step(next(iterator), model, metrics=metrics)
 
-  def test_task(self):
-    config = sentence_prediction.SentencePredictionConfig(
-        init_checkpoint=self.get_temp_dir(),
-        model=self.get_model_config(2),
-        train_data=self._train_data_config)
-    task = sentence_prediction.SentencePredictionTask(config)
-    model = task.build_model()
-    metrics = task.build_metrics()
-    dataset = task.build_inputs(config.train_data)
-
-    iterator = iter(dataset)
-    optimizer = tf.keras.optimizers.SGD(lr=0.1)
-    task.train_step(next(iterator), model, optimizer, metrics=metrics)
-    task.validation_step(next(iterator), model, metrics=metrics)
-
+  @parameterized.named_parameters(
+      ("init_cls_pooler", True),
+      ("init_encoder", False),
+  )
+  def test_task(self, init_cls_pooler):
     # Saves a checkpoint.
     pretrain_cfg = bert.PretrainerConfig(
         encoder=encoders.EncoderConfig(
             bert=encoders.BertEncoderConfig(vocab_size=30522, num_layers=1)),
         cls_heads=[
             bert.ClsHeadConfig(
-                inner_dim=10, num_classes=3, name="next_sentence")
+                inner_dim=768, num_classes=2, name="next_sentence")
         ])
     pretrain_model = masked_lm.MaskedLMTask(None).build_model(pretrain_cfg)
     ckpt = tf.train.Checkpoint(
         model=pretrain_model, **pretrain_model.checkpoint_items)
-    ckpt.save(config.init_checkpoint)
+    init_path = ckpt.save(self.get_temp_dir())
+
+    # Creates the task.
+    config = sentence_prediction.SentencePredictionConfig(
+        init_checkpoint=init_path,
+        model=self.get_model_config(num_classes=2),
+        train_data=self._train_data_config,
+        init_cls_pooler=init_cls_pooler)
+    task = sentence_prediction.SentencePredictionTask(config)
+    model = task.build_model()
+    metrics = task.build_metrics()
+    dataset = task.build_inputs(config.train_data)
+
+    iterator = iter(dataset)
+    optimizer = tf.keras.optimizers.SGD(lr=0.1)
     task.initialize(model)
+    task.train_step(next(iterator), model, optimizer, metrics=metrics)
+    task.validation_step(next(iterator), model, metrics=metrics)
 
   @parameterized.named_parameters(
       {