Add dropout module

bi_lstm_crf/app/train.py

@@ -99,6 +99,7 @@ def main():
     parser.add_argument('--num_epoch', type=int, default=20, help="number of epoch to train")
     parser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
     parser.add_argument('--weight_decay', type=float, default=0., help='the L2 normalization parameter')
+    parser.add_argument('--dropout', type=float, default=0., help='dropout rate for embedding, LSTM, and CRF')
     parser.add_argument('--batch_size', type=int, default=1000, help='batch size for training')
     parser.add_argument('--device', type=str, default=None,
                         help='the training device: "cuda:0", "cpu:0". It will be auto-detected by default')

bi_lstm_crf/app/utils.py

@@ -16,7 +16,7 @@ def model_filepath(model_dir):
 
 def build_model(args, processor, load=True, verbose=False):
     model = BiRnnCrf(len(processor.vocab), len(processor.tags),
-                     embedding_dim=args.embedding_dim, hidden_dim=args.hidden_dim, num_rnn_layers=args.num_rnn_layers)
+                     embedding_dim=args.embedding_dim, hidden_dim=args.hidden_dim, dropout=args.dropout, num_rnn_layers=args.num_rnn_layers)
 
     # weights
     model_path = model_filepath(args.model_dir)

bi_lstm_crf/model/model.py

@@ -4,22 +4,24 @@
 
 
 class BiRnnCrf(nn.Module):
-    def __init__(self, vocab_size, tagset_size, embedding_dim, hidden_dim, num_rnn_layers=1, rnn="lstm"):
+    def __init__(self, vocab_size, tagset_size, embedding_dim, hidden_dim, dropout=0.0, num_rnn_layers=1, rnn="lstm"):
         super(BiRnnCrf, self).__init__()
         self.embedding_dim = embedding_dim
         self.hidden_dim = hidden_dim
         self.vocab_size = vocab_size
         self.tagset_size = tagset_size
 
         self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.dropout = nn.Dropout(dropout)
         RNN = nn.LSTM if rnn == "lstm" else nn.GRU
         self.rnn = RNN(embedding_dim, hidden_dim // 2, num_layers=num_rnn_layers,
+                       dropout=dropout,
                        bidirectional=True, batch_first=True)
         self.crf = CRF(hidden_dim, self.tagset_size)
 
     def __build_features(self, sentences):
         masks = sentences.gt(0)
-        embeds = self.embedding(sentences.long())
+        embeds = self.dropout(self.embedding(sentences.long()))
 
         seq_length = masks.sum(1)
         sorted_seq_length, perm_idx = seq_length.sort(descending=True)
@@ -29,7 +31,7 @@ def __build_features(self, sentences):
         packed_output, _ = self.rnn(pack_sequence)
         lstm_out, _ = pad_packed_sequence(packed_output, batch_first=True)
         _, unperm_idx = perm_idx.sort()
-        lstm_out = lstm_out[unperm_idx, :]
+        lstm_out = self.dropout(lstm_out[unperm_idx, :])
 
         return lstm_out, masks