batch-level decoding in DuIE for memory issues (PaddlePaddle#253)

* fix duie out of memory when prediction
* batch-level decoding

Author: BenfengXu

examples/information_extraction/DuIE/data/id2spo.json

@@ -0,0 +1 @@
+{"predicate": ["empty", "empty", "注册资本", "作者", "所属专辑", "歌手", "邮政编码", "主演", "上映时间", "上映时间", "饰演", "饰演", "国籍", "成立日期", "毕业院校", "作曲", "作词", "编剧", "导演", "面积", "占地面积", "总部地点", "制片人", "嘉宾", "简称", "主持人", "获奖", "获奖", "获奖", "获奖", "海拔", "出品公司", "配音", "配音", "所在城市", "号", "主角", "创始人", "父亲", "祖籍", "母亲", "朝代", "董事长", "人口数量", "妻子", "丈夫", "票房", "票房", "专业代码", "气候", "修业年限", "改编自", "官方语言", "首都", "主题曲", "校长", "代言人"], "subject_type": ["empty", "empty", "企业", "图书作品", "歌曲", "歌曲", "行政区", "影视作品", "影视作品", "影视作品", "娱乐人物", "娱乐人物", "人物", "机构", "人物", "歌曲", "歌曲", "影视作品", "影视作品", "行政区", "机构", "企业", "影视作品", "电视综艺", "机构", "电视综艺", "娱乐人物", "娱乐人物", "娱乐人物", "娱乐人物", "地点", "影视作品", "娱乐人物", "娱乐人物", "景点", "历史人物", "文学作品", "企业", "人物", "人物", "人物", "历史人物", "企业", "行政区", "人物", "人物", "影视作品", "影视作品", "学科专业", "行政区", "学科专业", "影视作品", "国家", "国家", "影视作品", "学校", "企业/品牌"], "object_type": ["empty", "empty", "Number", "人物", "音乐专辑", "人物", "Text", "人物", "Date_@value", "地点_inArea", "人物_@value", "影视作品_inWork", "国家", "Date", "学校", "人物", "人物", "人物", "人物", "Number", "Number", "地点", "人物", "人物", "Text", "人物", "奖项_@value", "作品_inWork", "Date_onDate", "Number_period", "Number", "企业", "人物_@value", "影视作品_inWork", "城市", "Text", "人物", "人物", "人物", "地点", "人物", "Text", "人物", "Number", "人物", "人物", "Number_@value", "地点_inArea", "Text", "气候", "Number", "作品", "语言", "城市", "歌曲", "人物", "人物"]}

examples/information_extraction/DuIE/data/predicate2id.json

@@ -0,0 +1 @@
+{"O": 0, "I": 1, "注册资本": 2, "作者": 3, "所属专辑": 4, "歌手": 5, "邮政编码": 6, "主演": 7, "上映时间_@value": 8, "上映时间_inArea": 9, "饰演_@value": 10, "饰演_inWork": 11, "国籍": 12, "成立日期": 13, "毕业院校": 14, "作曲": 15, "作词": 16, "编剧": 17, "导演": 18, "面积": 19, "占地面积": 20, "总部地点": 21, "制片人": 22, "嘉宾": 23, "简称": 24, "主持人": 25, "获奖_@value": 26, "获奖_inWork": 27, "获奖_onDate": 28, "获奖_period": 29, "海拔": 30, "出品公司": 31, "配音_@value": 32, "配音_inWork": 33, "所在城市": 34, "号": 35, "主角": 36, "创始人": 37, "父亲": 38, "祖籍": 39, "母亲": 40, "朝代": 41, "董事长": 42, "人口数量": 43, "妻子": 44, "丈夫": 45, "票房_@value": 46, "票房_inArea": 47, "专业代码": 48, "气候": 49, "修业年限": 50, "改编自": 51, "官方语言": 52, "首都": 53, "主题曲": 54, "校长": 55, "代言人": 56}

examples/information_extraction/DuIE/run_duie.py

@@ -88,47 +88,41 @@ def evaluate(model, criterion, data_loader, file_path, mode):
         predict_test.json and predict_test.json.zip \
         under args.data_path dir for later submission or evaluation.
     """
+    example_all = []
+    with open(file_path, "r", encoding="utf-8") as fp:
+        for line in fp:
+            example_all.append(json.loads(line))
+    id2spo_path = os.path.join(os.path.dirname(file_path), "id2spo.json")
+    with open(id2spo_path, 'r', encoding='utf8') as fp:
+        id2spo = json.load(fp)
+
     model.eval()
-    probs_all = None
-    seq_len_all = None
-    tok_to_orig_start_index_all = None
-    tok_to_orig_end_index_all = None
     loss_all = 0
     eval_steps = 0
+    formatted_outputs = []
+    current_idx = 0
     for batch in tqdm(data_loader, total=len(data_loader)):
         eval_steps += 1
         input_ids, seq_len, tok_to_orig_start_index, tok_to_orig_end_index, labels = batch
         logits = model(input_ids=input_ids)
-        mask = (input_ids != 0).logical_and((input_ids != 1)).logical_and(
-            (input_ids != 2))
+        mask = (input_ids != 0).logical_and((input_ids != 1)).logical_and((input_ids != 2))
         loss = criterion(logits, labels, mask)
         loss_all += loss.numpy().item()
         probs = F.sigmoid(logits)
-        if probs_all is None:
-            probs_all = probs.numpy()
-            seq_len_all = seq_len.numpy()
-            tok_to_orig_start_index_all = tok_to_orig_start_index.numpy()
-            tok_to_orig_end_index_all = tok_to_orig_end_index.numpy()
-        else:
-            probs_all = np.append(probs_all, probs.numpy(), axis=0)
-            seq_len_all = np.append(seq_len_all, seq_len.numpy(), axis=0)
-            tok_to_orig_start_index_all = np.append(
-                tok_to_orig_start_index_all,
-                tok_to_orig_start_index.numpy(),
-                axis=0)
-            tok_to_orig_end_index_all = np.append(
-                tok_to_orig_end_index_all,
-                tok_to_orig_end_index.numpy(),
-                axis=0)
+        logits_batch = probs.numpy()
+        seq_len_batch = seq_len.numpy()
+        tok_to_orig_start_index_batch = tok_to_orig_start_index.numpy()
+        tok_to_orig_end_index_batch = tok_to_orig_end_index.numpy()
+        formatted_outputs.extend(decoding(example_all[current_idx: current_idx+len(logits)],
+                                          id2spo,
+                                          logits_batch,
+                                          seq_len_batch,
+                                          tok_to_orig_start_index_batch,
+                                          tok_to_orig_end_index_batch))
+        current_idx = current_idx+len(logits)
     loss_avg = loss_all / eval_steps
     print("eval loss: %f" % (loss_avg))
 
-    id2spo_path = os.path.join(os.path.dirname(file_path), "id2spo.json")
-    with open(id2spo_path, 'r', encoding='utf8') as fp:
-        id2spo = json.load(fp)
-    formatted_outputs = decoding(file_path, id2spo, probs_all, seq_len_all,
-                                 tok_to_orig_start_index_all,
-                                 tok_to_orig_end_index_all)
     if mode == "predict":
         predict_file_path = os.path.join(args.data_path, 'predictions.json')
     else:
@@ -228,7 +222,7 @@ def do_train():
             optimizer.clear_grad()
             loss_item = loss.numpy().item()
             global_step += 1
-
+            
             if global_step % logging_steps == 0 and rank == 0:
                 print(
                     "epoch: %d / %d, steps: %d / %d, loss: %f, speed: %.2f step/s"
@@ -270,7 +264,7 @@ def do_train():
 
 def do_predict():
     paddle.set_device(args.device)
-
+    
     # Reads label_map.
     label_map_path = os.path.join(args.data_path, "predicate2id.json")
     if not (os.path.exists(label_map_path) and os.path.isfile(label_map_path)):
@@ -313,6 +307,7 @@ def do_predict():
 
 
 if __name__ == "__main__":
+
     if args.do_train:
         do_train()
     elif args.do_predict:

examples/information_extraction/DuIE/utils.py

@@ -42,19 +42,18 @@ def find_entity(text_raw, id_, predictions, tok_to_orig_start_index,
     return list(set(entity_list))
 
 
-def decoding(file_path, id2spo, logits_all, seq_len_all,
-             tok_to_orig_start_index_all, tok_to_orig_end_index_all):
+def decoding(example_batch,
+             id2spo,
+             logits_batch,
+             seq_len_batch,
+             tok_to_orig_start_index_batch,
+             tok_to_orig_end_index_batch):
     """
     model output logits -> formatted spo (as in data set file)
     """
-    example_all = []
-    with open(file_path, "r", encoding="utf-8") as fp:
-        for line in fp:
-            example_all.append(json.loads(line))
-
     formatted_outputs = []
     for (i, (example, logits, seq_len, tok_to_orig_start_index, tok_to_orig_end_index)) in \
-            enumerate(zip(example_all, logits_all, seq_len_all, tok_to_orig_start_index_all, tok_to_orig_end_index_all)):
+            enumerate(zip(example_batch, logits_batch, seq_len_batch, tok_to_orig_start_index_batch, tok_to_orig_end_index_batch)):
 
         logits = logits[1:seq_len +
                         1]  # slice between [CLS] and [SEP] to get valid logits