Fixed bug when NER weight not being updated

Author: LuciusLan

main.py

@@ -153,7 +153,7 @@ def custom_batch_collation(x):
 seg_loss = FocalLoss(ignore_index=0, gamma=2)
 
 bert_param_optimizer = list(model.transformer.named_parameters())
-ner_fc_param_optimizer = list(model.plain_ner.named_parameters())
+ner_fc_param_optimizer = list(model.ner.named_parameters())
 no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
 
 if not BASELINE:

model.py

@@ -19,8 +19,9 @@ def __init__(self, config):
                 pretrained_model_name_or_path="roberta-base", cache_dir=MODEL_CACHE_DIR, config=config)
         self.dropout = nn.Dropout()
         self.relu = nn.ReLU(True)
-        self.plain_ner = nn.Linear(config.hidden_size, len(LABEL_BIO))
-        if not BASELINE:
+        if BASELINE:
+            self.ner = nn.Linear(config.hidden_size, len(LABEL_BIO))
+        else:
             self.boundary_encoder = nn.LSTM(bidirectional=True, input_size=config.hidden_size, hidden_size=LSTM_HIDDEN, batch_first=True)
             self.boundary_decoder = nn.LSTM(bidirectional=False, input_size=LSTM_HIDDEN*2, hidden_size=LSTM_HIDDEN, batch_first=True)
             self.boundary_biaffine = BoundaryBiaffine(LSTM_HIDDEN, LSTM_HIDDEN*2, len(BOUNDARY_LABEL_UNIDIRECTION))