fix ofa export bug (PaddlePaddle#1322)

Author: LiuChiachi

examples/model_compression/ofa/export_model.py

@@ -27,12 +27,35 @@
 import paddle.nn.functional as F
 
 from paddlenlp.transformers import BertModel, BertForSequenceClassification, BertTokenizer
+from paddlenlp.transformers import TinyBertModel, TinyBertForSequenceClassification, TinyBertTokenizer
+from paddlenlp.transformers import TinyBertForSequenceClassification, TinyBertTokenizer
+from paddlenlp.transformers import RobertaForSequenceClassification, RobertaTokenizer
 from paddlenlp.utils.log import logger
 from paddleslim.nas.ofa import OFA, utils
 from paddleslim.nas.ofa.convert_super import Convert, supernet
 from paddleslim.nas.ofa.layers import BaseBlock
 
-MODEL_CLASSES = {"bert": (BertForSequenceClassification, BertTokenizer), }
+MODEL_CLASSES = {
+    "bert": (BertForSequenceClassification, BertTokenizer),
+    "roberta": (RobertaForSequenceClassification, RobertaTokenizer),
+    "tinybert": (TinyBertForSequenceClassification, TinyBertTokenizer),
+}
+
+
+def tinybert_forward(self, input_ids, token_type_ids=None, attention_mask=None):
+    wtype = self.pooler.dense.fn.weight.dtype if hasattr(
+        self.pooler.dense, 'fn') else self.pooler.dense.weight.dtype
+    if attention_mask is None:
+        attention_mask = paddle.unsqueeze(
+            (input_ids == self.pad_token_id).astype(wtype) * -1e9, axis=[1, 2])
+    embedding_output = self.embeddings(input_ids, token_type_ids)
+    encoded_layer = self.encoder(embedding_output, attention_mask)
+    pooled_output = self.pooler(encoded_layer)
+
+    return encoded_layer, pooled_output
+
+
+TinyBertModel.forward = tinybert_forward
 
 
 def parse_args():
@@ -113,14 +136,15 @@ def do_train(args):
     config_path = os.path.join(args.model_name_or_path, 'model_config.json')
     cfg_dict = dict(json.loads(open(config_path).read()))
 
+    kept_layers_index = {}
     if args.depth_mult < 1.0:
-        depth = round(cfg_dict["init_args"][0]['num_hidden_layers'] * args.depth_mult)
-        cfg_dict["init_args"][0]['num_hidden_layers'] = depth 
-        kept_layers_index = {}
-        for idx, i in enumerate(range(1, depth+1)):
+        depth = round(cfg_dict["init_args"][0]['num_hidden_layers'] *
+                      args.depth_mult)
+        cfg_dict["init_args"][0]['num_hidden_layers'] = depth
+        for idx, i in enumerate(range(1, depth + 1)):
             kept_layers_index[idx] = math.floor(i / args.depth_mult) - 1
 
-    os.rename(config_path, config_path+'_bak')
+    os.rename(config_path, config_path + '_bak')
     with open(config_path, "w", encoding="utf-8") as f:
         f.write(json.dumps(cfg_dict, ensure_ascii=False))
 
@@ -132,7 +156,7 @@ def do_train(args):
     origin_model = model_class.from_pretrained(
         args.model_name_or_path, num_classes=num_labels)
 
-    os.rename(config_path+'_bak', config_path)
+    os.rename(config_path + '_bak', config_path)
 
     sp_config = supernet(expand_ratio=[1.0, args.width_mult])
     model = Convert(sp_config).convert(model)
@@ -142,15 +166,24 @@ def do_train(args):
     sd = paddle.load(
         os.path.join(args.model_name_or_path, 'model_state.pdparams'))
 
-    for name, params in ofa_model.model.named_parameters():
-        if 'encoder' not in name:
-            params.set_value(sd[name])
-        else:
-            idx = int(name.strip().split('.')[3])
-            mapping_name = name.replace('.'+str(idx)+'.', '.'+str(kept_layers_index[idx])+'.')
-            params.set_value(sd[mapping_name])
+    if len(kept_layers_index) == 0:
+        ofa_model.model.set_state_dict(sd)
+    else:
+        for name, params in ofa_model.model.named_parameters():
+            if 'encoder' not in name:
+                params.set_value(sd[name])
+            else:
+                idx = int(name.strip().split('.')[3])
+                mapping_name = name.replace(
+                    '.' + str(idx) + '.',
+                    '.' + str(kept_layers_index[idx]) + '.')
+                params.set_value(sd[mapping_name])
 
     best_config = utils.dynabert_config(ofa_model, args.width_mult)
+    for name, sublayer in ofa_model.model.named_sublayers():
+        if isinstance(sublayer, paddle.nn.MultiHeadAttention):
+            sublayer.num_heads = int(args.width_mult * sublayer.num_heads)
+
     ofa_model.export(
         best_config,
         input_shapes=[[1, args.max_seq_length], [1, args.max_seq_length]],