update lm config

README.md

@@ -81,7 +81,7 @@ sentence|relation|head|head_type|head_offset|tail|tail_type|tail_offset
 
 使用语言预训练模型时，要提前下载好预训练好的参数，放到 `pretrained` 文件夹内。
 
-另外数据量很小时，直接使用如12层的bert，效果并不理想，反而层数调低些收敛更快效果更好。可以将 bert `config.json` 文件中下 `num_hidden_layers` 参数调整到如3，6等层数。具体看训练的数据量大小。
+另外数据量较小时，直接使用如12层的bert，效果并不理想，反而层数调低些收敛更快效果更好。如调整到3、6层数，具体看训练的数据量大小。
 
 ## 后续工作
 

deepke/config.py

@@ -53,8 +53,11 @@ class CapsuleConfig(object):
 
 class LMConfig(object):
     # lm_name = 'bert-base-chinese'  # download usage
-    lm_file = 'bert_pretrained'  # cache file usage
-
+    # cache file usage
+    lm_file = 'bert_pretrained'
+    # transformer 层数，初始 base bert 为12层
+    # 但是数据量较小时调低些反而收敛更快效果更好
+    num_hidden_layers = 2
 
 class Config(object):
     # 原始数据存放位置

deepke/model/LM.py

@@ -10,7 +10,7 @@ class LM(BasicModule):
         self.lm_name = config.lm.lm_file
         self.out_dim = config.relation_type
 
-        self.lm = BertModel.from_pretrained(self.lm_name)
+        self.lm = BertModel.from_pretrained(self.lm_name, num_hidden_layers=config.lm.num_hidden_layers)
         self.fc = nn.Linear(768, self.out_dim)
 
     def forward(self, x):