add hidden size

Author: FreeRotate

models/LSTM.py

@@ -12,22 +12,18 @@ def __init__(self, vocab_len, config):
         super().__init__()
         self.num_labels = config.num_labels
         self.embed = nn.Embedding(num_embeddings=vocab_len, embedding_dim=config.embed_dim)
-        self.dropout = nn.Dropout(config.dropout_rate)
-        self.lstm = nn.LSTM(input_size=config.embed_dim, hidden_size=config.embed_dim, batch_first=True, bidirectional=True)
-        self.layer_norm = nn.LayerNorm(config.embed_dim * 2)
-        self.classifier = nn.Linear(config.embed_dim * 2, config.num_labels)
+        self.lstm = nn.LSTM(input_size=config.embed_dim, hidden_size=config.hidden_size, batch_first=True, bidirectional=True)
+        self.classifier = nn.Linear(config.hidden_size * 2, config.num_labels)
         self.crf = CRF(num_tags=config.num_labels, batch_first=True)
         self.loss_fct = nn.CrossEntropyLoss()
 
     def forward(self, word_ids, label_ids=None, label_mask=None, use_crf=True):
         word_embed = self.embed(word_ids)
-        word_embed = self.dropout(word_embed)
         sequence_output, _ = self.lstm(word_embed)
-        sequence_output = self.layer_norm(sequence_output)
         logits = self.classifier(sequence_output)
         if label_ids != None:
             if use_crf:
-                loss = self.crf(logits, label_ids)
+                loss = self.crf(logits, label_ids, label_mask)
                 loss = -1 * loss
             else:
                 active_logits = logits.view(-1, self.num_labels)