try 内存回收

PointerNet/main.py

@@ -34,7 +34,8 @@ class NerPipeline:
         torch.save(self.optimizer.state_dict(), self.args.optimizer_save_dir)
 
     def load_model(self):
-        self.model.load_state_dict(torch.load(self.args.save_dir, map_location="cpu"))
+        # self.model.load_state_dict(torch.load(self.args.save_dir, map_location="cpu"))  #GPU 上训练的模型加载到CPU
+        self.model.load_state_dict(torch.load(self.args.save_dir))
         self.model.to(self.args.device)
 
     def build_optimizer_and_scheduler(self, t_total):

PointerNet/model.py

@@ -1,4 +1,5 @@
 import torch
+import gc
 import torch.nn as nn
 from transformers import BertConfig, BertModel, BertForSequenceClassification
 
@@ -165,6 +166,9 @@ class UIEModel(nn.Module):
         cls_bert_output = outputs.pooler_output  # [sent_num, hidden_size] 
         # dropout
         pooled_output = self.dropout(cls_bert_output)
+        # 尝试内存回收
+        del outputs,cls_bert_output
+        gc.collect()
 
         # 对每个pointer位置接个线性层
         # 对试题判断任务也接个线性层