crf implementation

atulkum · atulkum · commit 110bfce36a4b · 2019-01-11T08:25:47.000-08:00
diff --git a/neural_ner/config.py b/neural_ner/config.py
@@ -42,6 +42,4 @@ class Config(object):
 
 config.is_cuda = False
 
-config.is_l2_loss = False
-
-config.optimizer='sdg' #'adam'
+config.is_l2_loss = False
diff --git a/neural_ner/model.py b/neural_ner/model.py
@@ -8,6 +8,8 @@
 import logging
 import numpy as np
 
+from crf import CRF_Loss
+
 print('pytorch version', torch.__version__)
 
 logging.basicConfig(level=logging.INFO)
@@ -58,7 +60,8 @@ def get_word_embd(vocab, config):
 def test_one_batch(batch, model):
     model.eval()
     logits = model(batch)
-    _, pred = model.get_argmax(logits)
+    lengths = batch['words_lens']
+    pred = model.predict(logits, lengths)
     return logits, pred
 
 def get_model(vocab, config, model_file_path, is_eval=False):
@@ -80,7 +83,6 @@ def __init__(self, vocab, config):
         super(NER_SOFTMAX_CHAR, self).__init__()
         word_emb_matrix = get_word_embd(vocab, config)
         embd_vector = torch.from_numpy(word_emb_matrix).float()
-        tagset_size = len(vocab.id_to_tag)
 
         self.word_embeds = nn.Embedding.from_pretrained(embd_vector, freeze=False)
         self.char_embeds = nn.Embedding(len(vocab.char_to_id), config.char_embd_dim, padding_idx=Constants.PAD_ID)
@@ -97,7 +99,8 @@ def __init__(self, vocab, config):
         self.dropout = nn.Dropout(config.dropout_rate)
         self.hidden_layer = nn.Linear(config.word_lstm_dim * 2, config.word_lstm_dim)
         self.tanh_layer = torch.nn.Tanh()
-        self.hidden2tag = nn.Linear(config.word_lstm_dim, tagset_size)
+
+        self.hidden2tag = nn.Linear(config.word_lstm_dim, len(vocab.id_to_tag))
 
         self.config = config
 
@@ -158,144 +161,50 @@ def forward(self, batch):
 
         return logits
 
-    def neg_log_likelihood(self, logits, y, s_len):
+    def neg_log_likelihood(self, logits, y, s_lens):
         log_smx = F.log_softmax(logits, dim=2)
         loss = F.nll_loss(log_smx.transpose(1, 2), y, ignore_index=Constants.TAG_PAD_ID, reduce=False)
-        loss = loss.squeeze(1).sum(dim=1) / s_len.float()
+        loss = loss.squeeze(1).sum(dim=1) / s_lens.float()
         loss = loss.mean()
-        if self.config.is_l2_loss:
-            l2_reg = sum(p.norm(2) for p in self.parameters() if p.requires_grad)
-            loss += self.config.reg_lambda * l2_reg
         return loss
 
-    def get_argmax(self, logits):
-        max_value, max_idx = torch.max(logits, dim=2)
-        return max_value, max_idx
-
-class NER_CRF(nn.Module):
-    def __init__(self, embd_vector, hidden_dim, tagset_size,
-                 reg_lambda):
-        super(NER_CRF, self).__init__()
-
-        self.start_tag_idx = tagset_size
-        self.stop_tag_idx = tagset_size + 1
-        self.all_tagset_size = tagset_size + 2
-
-        self.word_embeds = nn.Embedding.from_pretrained(embd_vector)
-        embedding_dim = self.word_embeds.embedding_dim
-
-        self.lstm = nn.LSTM(embedding_dim, hidden_dim//2,
-                            num_layers=1, bidirectional=True, batch_first=True)
-        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)
-
-        #transition from y_i-1 to y_i, T[y_i, y_j] = y_i <= y_j
-        #+2 added for start and end indices
-        self.transitions = nn.Parameter(torch.randn(self.all_tagset_size, self.all_tagset_size))
-        nn.init.uniform_(self.transitions, -0.1, 0.1)
-
-        #no transition to start_tag, not transition from end tag
-        self.transitions.data[self.start_tag_idx, :] = -10000
-        self.transitions.data[:, self.stop_tag_idx] = -10000
-
-        self.hidden = self.init_hidden()
-
-    def init_hidden(self):
-        return (torch.randn(2, 1, self.hidden_dim // 2),
-                torch.randn(2, 1, self.hidden_dim // 2))
+    def get_loss(self, logits, y, s_lens):
+        loss = self.neg_log_likelihood(logits, y, s_lens)
+        if self.config.is_l2_loss:
+            loss += self.get_l2_loss()
+        return loss
 
-    def get_emission_prob(self, sentence, lengths):
-        embedded = self.word_embeds(sentence)
-        lengths = lengths.view(-1).tolist()
-        packed = pack_padded_sequence(embedded, lengths, batch_first=True)
+    def get_l2_loss(self):
+        l2_reg = sum(p.norm(2) for p in self.parameters() if p.requires_grad)
+        return self.config.reg_lambda * l2_reg
 
-        self.hidden = self.init_hidden()
-        output, self.hidden = self.lstm(packed, self.hidden)
+    def predict(self, logit, lengths):
+        max_value, pred = torch.max(logit, dim=2)
+        return pred
 
-        lstm_feats, _ = pad_packed_sequence(output, batch_first=True)  # h dim = B x t_k x n
-        lstm_feats = lstm_feats.contiguous()
+class NER_SOFTMAX_CHAR_CRF(nn.Module):
+    def __init__(self, vocab, config):
+        super(NER_SOFTMAX_CHAR_CRF, self).__init__()
 
-        b, t_k, d = list(lstm_feats.size())
+        self.featurizer = NER_SOFTMAX_CHAR(vocab, config)
+        self.crf = CRF_Loss(len(vocab.id_to_tag))
 
-        logits = self.hidden2tag(lstm_feats.view(-1, d))
-        logits = logits.view(b, t_k, -1)
+    def forward(self, batch):
+        emissions = self.featurizer(batch)
+        return emissions
 
-        return logits
+    def crf_loss(self, emissions, target, s_lens):
+        loss = -1 * self.crf(emissions, target)
+        loss = loss.squeeze(1).sum(dim=1) / s_lens.float()
+        loss = loss.mean()
+        return loss
 
-    def get_argmax(self, logits):
-        max_value, max_idx = torch.max(logits, dim=2)
-        return max_value, max_idx
-
-    def log_sum_exp(self, vec):
-        max_score, _ = self.get_argmax(vec)
-        max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
-        return max_score + torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))
-
-    def get_log_z(self, emission_prob):
-        init_alphas = torch.full((1, self.all_tagset_size), -10000.)
-        init_alphas[0][self.start_tag_idx] = 0.
-        forward_var = init_alphas
-        for e_i in emission_prob:
-            alphas_t = []
-            for next_tag in range(self.all_tagset_size):
-                emit_score = e_i[next_tag].view(
-                    1, -1).expand(1, self.all_tagset_size)
-                trans_score = self.transitions[next_tag].view(1, -1)
-
-                next_tag_var = forward_var + trans_score + emit_score
-                alphas_t.append(self.log_sum_exp(next_tag_var).view(1))
-            forward_var = torch.cat(alphas_t).view(1, -1)
-        terminal_var = forward_var + self.transitions[self.stop_tag_idx]
-        alpha = self.log_sum_exp(terminal_var)
-        return alpha
-
-    def get_log_p_y_x(self, feats, lengths, tags):
-        score = torch.zeros(1)
-        tags = torch.cat([torch.tensor([self.start_tag_idx], dtype=torch.long), tags])
-        for i, feat in enumerate(feats):
-            score = score + self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
-        score = score + self.transitions[self.stop_tag_idx, tags[-1]]
-        return score
-
-    def neg_log_likelihood(self, sentence, lengths, tags):
-        feats = self.get_emission_prob(sentence, lengths)
-        log_z = self.get_log_z(feats, lengths)
-        log_p_y_x = self.get_log_p_y_x(feats, tags, lengths)
-        return -(log_p_y_x - log_z)
-
-    def forward(self, sentence, lengths):
-        feats = self.get_emission_prob(sentence, lengths)
-        score, tag_seq = self.viterbi_decode(feats, lengths)
-        return score, tag_seq
-
-    def viterbi_decode(self, feats, lengths):
-        backpointers = []
-
-        init_vvars = torch.full((1, self.all_tagset_size), -10000.)
-        init_vvars[0][self.start_tag_idx] = 0
-
-        forward_var = init_vvars
-        for feat in feats:
-            bptrs_t = []
-            viterbivars_t = []
-
-            for next_tag in range(self.all_tagset_size):
-                next_tag_var = forward_var + self.transitions[next_tag]
-                _, best_tag_id = self.get_argmax(next_tag_var)
-                bptrs_t.append(best_tag_id)
-                viterbivars_t.append(next_tag_var[0][best_tag_id].view(1))
-            forward_var = (torch.cat(viterbivars_t) + feat).view(1, -1)
-            backpointers.append(bptrs_t)
-
-        terminal_var = forward_var + self.transitions[self.stop_tag_idx]
-        _, best_tag_id = self.get_argmax(terminal_var)
-        path_score = terminal_var[0][best_tag_id]
-
-        best_path = [best_tag_id]
-        for bptrs_t in reversed(backpointers):
-            best_tag_id = bptrs_t[best_tag_id]
-            best_path.append(best_tag_id)
-        start = best_path.pop()
-        assert start == self.start_tag_idx
-        best_path.reverse()
-        return path_score, best_path
+    def get_loss(self, logits, y, s_lens):
+        loss = self.crf_loss(logits, y, s_lens)
+        if self.config.is_l2_loss:
+            loss += self.get_l2_loss()
+        return loss
 
+    def predict(self, emissions, lengths):
+        best_scores, pred = self.crf.viterbi_decode(emissions, lengths)
+        return pred
diff --git a/neural_ner/process_training.py b/neural_ner/process_training.py
@@ -9,8 +9,7 @@
 
 from data_utils.batcher import DatasetConll2003
 from data_utils.vocab import Vocab
-from model import get_model
-from model import test_one_batch
+from model import get_model, test_one_batch
 from train_utils import setup_train_dir, save_model, write_summary, \
     get_param_norm, get_grad_norm, Evaluter
 
@@ -29,7 +28,7 @@ def train_one_batch(self, batch, optimizer, params):
         s_lengths = batch['words_lens']
         y = batch['tags']
         logits = self.model(batch)
-        loss = self.model.neg_log_likelihood(logits, y, s_lengths)
+        loss = self.model.get_loss(logits, y, s_lengths)
 
         loss.backward()
 
@@ -45,10 +44,7 @@ def train(self):
         train_dir, summary_writer = setup_train_dir(self.config)
 
         params = list(filter(lambda p: p.requires_grad, self.model.parameters()))
-        if self.config.optimizer == 'adam':
-            optimizer = Adam(params, lr=0.001, amsgrad=True)
-        elif self.config.optimizer == 'sdg':
-            optimizer = SGD(params, lr=0.01)
+        optimizer = Adam(params, lr=0.001, amsgrad=True)
 
         num_params = sum(p.numel() for p in params)
         logging.info("Number of params: %d" % num_params)
@@ -131,7 +127,7 @@ def evaluate(self, data_type, num_samples=None):
             y = batch['tags']
 
             logits, pred = test_one_batch(batch, self.model)
-            loss = self.model.neg_log_likelihood(logits, y, s_lengths)
+            loss = self.model.get_loss(logits, y, s_lengths)
 
             curr_batch_size = len(batch['raw_sentence'])
             loss_per_batch += loss * curr_batch_size
