Making the PretrainedMixin work for both encoder and classifier

Author: Atif Ahmed

texar/torch/modules/classifiers/bert_classifier.py

@@ -71,6 +71,8 @@ def __init__(self,
 
         super().__init__(hparams=hparams)
 
+        self.load_pretrained_config(pretrained_model_name, cache_dir)
+
         # Create the underlying encoder
         encoder_hparams = dict_fetch(hparams,
                                      self._ENCODER_CLASS.default_hparams())
@@ -120,6 +122,8 @@ def __init__(self,
                          (self.num_classes <= 0 and
                           self._hparams.encoder.dim == 1)
 
+        self.init_pretrained_weights(class_type='classifier')
+
     @staticmethod
     def default_hparams():
         r"""Returns a dictionary of hyperparameters with default values.

texar/torch/modules/pretrained/bert.py

@@ -326,12 +326,20 @@ def _init_from_checkpoint(self, pretrained_model_name: str,
         pooler_map = {
             'bert/pooler/dense/bias': 'pooler.0.bias',
             'bert/pooler/dense/kernel': 'pooler.0.weight',
+        }
+        classifier_map = {
             'output_bias': '_logits_layer.bias',
             'output_weights': '_logits_layer.weight',
         }
+        global_prefix_map = {
+            'classifier': '_encoder.'
+        }
         tf_path = os.path.abspath(os.path.join(
             cache_dir, self._MODEL2CKPT[pretrained_model_name]))
 
+        class_type = kwargs.get('class_type', 'encoder')
+        global_prefix = global_prefix_map.get(class_type, '')
+
         # Load weights from TF model
         init_vars = tf.train.list_variables(tf_path)
         tfnames, arrays = [], []
@@ -351,13 +359,14 @@ def _init_from_checkpoint(self, pretrained_model_name: str,
                 continue
 
             if name in global_tensor_map:
-                v_name = global_tensor_map[name]
+                v_name = global_prefix + global_tensor_map[name]
                 pointer = self._name_to_variable(v_name)
                 assert pointer.shape == array.shape
                 pointer.data = torch.from_numpy(array)
                 idx += 1
             elif name in pooler_map:
-                pointer = self._name_to_variable(pooler_map[name])
+                pointer = self._name_to_variable(global_prefix +
+                                                 pooler_map[name])
                 if name.endswith('bias'):
                     assert pointer.shape == array.shape
                     pointer.data = torch.from_numpy(array)
@@ -367,6 +376,13 @@ def _init_from_checkpoint(self, pretrained_model_name: str,
                     assert pointer.shape == array_t.shape
                     pointer.data = torch.from_numpy(array_t)
                     idx += 1
+            elif name in classifier_map:
+                if class_type != 'classifier':
+                    continue
+                pointer = self._name_to_variable(classifier_map[name])
+                assert pointer.shape == array.shape
+                pointer.data = torch.from_numpy(array)
+                idx += 1
             else:
                 # here name is the TensorFlow variable name
                 name_tmp = name.split("/")
@@ -375,12 +391,14 @@ def _init_from_checkpoint(self, pretrained_model_name: str,
                 name_tmp = "/".join(name_tmp[3:])
                 if name_tmp in layer_tensor_map:
                     v_name = layer_tensor_map[name_tmp].format(layer_no)
-                    pointer = self._name_to_variable(py_prefix + v_name)
+                    pointer = self._name_to_variable(global_prefix +
+                                                     py_prefix + v_name)
                     assert pointer.shape == array.shape
                     pointer.data = torch.from_numpy(array)
                 elif name_tmp in layer_transpose_map:
                     v_name = layer_transpose_map[name_tmp].format(layer_no)
-                    pointer = self._name_to_variable(py_prefix + v_name)
+                    pointer = self._name_to_variable(global_prefix +
+                                                     py_prefix + v_name)
                     array_t = np.transpose(array)
                     assert pointer.shape == array_t.shape
                     pointer.data = torch.from_numpy(array_t)