├── LSTM_CRF_faster.py
├── LSTM_CRF_faster_parallel.py
└── README.md


/LSTM_CRF_faster.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | import torch.optim as optim
  4 | START_TAG = "<START>"
  5 | STOP_TAG = "<STOP>"
  6 | EMBEDDING_DIM = 5
  7 | HIDDEN_DIM = 4
  8 | # torch.manual_seed(1)
  9 | 
 10 | def prepare_sequence(seq, to_ix):
 11 |     idxs = [to_ix[w] for w in seq]
 12 |     return torch.tensor(idxs, dtype=torch.long)
 13 | 
 14 | class BiLSTM_CRF(nn.Module):
 15 | 
 16 |     def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):
 17 |         super(BiLSTM_CRF, self).__init__()
 18 |         self.embedding_dim = embedding_dim
 19 |         self.hidden_dim = hidden_dim
 20 |         self.vocab_size = vocab_size
 21 |         self.tag_to_ix = tag_to_ix
 22 |         self.tagset_size = len(tag_to_ix)
 23 | 
 24 |         self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
 25 |         self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2,
 26 |                             num_layers=1, bidirectional=True)
 27 | 
 28 |         # Maps the output of the LSTM into tag space.
 29 |         self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
 30 | 
 31 |         # Matrix of transition parameters.  Entry i,j is the score of
 32 |         # transitioning *to* i *from* j.
 33 |         self.transitions = nn.Parameter(
 34 |             torch.randn(self.tagset_size, self.tagset_size))
 35 | 
 36 |         # These two statements enforce the constraint that we never transfer
 37 |         # to the start tag and we never transfer from the stop tag
 38 |         self.transitions.data[tag_to_ix[START_TAG], :] = -10000
 39 |         self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000
 40 | 
 41 |         self.hidden = self.init_hidden()
 42 | 
 43 |     def init_hidden(self):
 44 |         return (torch.randn(2, 1, self.hidden_dim // 2),
 45 |                 torch.randn(2, 1, self.hidden_dim // 2))
 46 | 
 47 |     def _forward_alg(self, feats):
 48 |         # Do the forward algorithm to compute the partition function
 49 |         init_alphas = torch.full([self.tagset_size], -10000.)
 50 |         # START_TAG has all of the score.
 51 |         init_alphas[self.tag_to_ix[START_TAG]] = 0.
 52 | 
 53 |         # Wrap in a variable so that we will get automatic backprop
 54 |         # Iterate through the sentence
 55 |         forward_var_list=[]
 56 |         forward_var_list.append(init_alphas)
 57 |         for feat_index in range(feats.shape[0]):
 58 |             gamar_r_l = torch.stack([forward_var_list[feat_index]] * feats.shape[1])
 59 |             t_r1_k = torch.unsqueeze(feats[feat_index],0).transpose(0,1)
 60 |             aa = gamar_r_l + t_r1_k + self.transitions
 61 |             forward_var_list.append(torch.logsumexp(aa,dim=1))
 62 |         terminal_var = forward_var_list[-1] + self.transitions[self.tag_to_ix[STOP_TAG]]
 63 |         terminal_var = torch.unsqueeze(terminal_var,0)
 64 |         alpha = torch.logsumexp(terminal_var, dim=1)[0]
 65 |         return alpha
 66 | 
 67 |     def _get_lstm_features(self, sentence):
 68 |         self.hidden = self.init_hidden()
 69 |         embeds = self.word_embeds(sentence).view(len(sentence), 1, -1)
 70 |         lstm_out, self.hidden = self.lstm(embeds, self.hidden)
 71 |         lstm_out = lstm_out.view(len(sentence), self.hidden_dim)
 72 |         lstm_feats = self.hidden2tag(lstm_out)
 73 |         return lstm_feats
 74 | 
 75 |     def _score_sentence(self, feats, tags):
 76 |         # Gives the score of a provided tag sequence
 77 |         score = torch.zeros(1)
 78 |         tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])
 79 |         for i, feat in enumerate(feats):
 80 |             score = score + \
 81 |                 self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
 82 |         score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
 83 |         return score
 84 | 
 85 |     def _viterbi_decode(self, feats):
 86 |         backpointers = []
 87 |         # Initialize the viterbi variables in log space
 88 |         init_vvars = torch.full((1, self.tagset_size), -10000.)
 89 |         init_vvars[0][self.tag_to_ix[START_TAG]] = 0
 90 | 
 91 |         # forward_var at step i holds the viterbi variables for step i-1
 92 |         forward_var_list = []
 93 |         forward_var_list.append(init_vvars)
 94 | 
 95 |         for feat_index in range(feats.shape[0]):
 96 |             gamar_r_l = torch.stack([forward_var_list[feat_index]] * feats.shape[1])
 97 |             gamar_r_l = torch.squeeze(gamar_r_l)
 98 |             next_tag_var = gamar_r_l + self.transitions
 99 |             viterbivars_t,bptrs_t = torch.max(next_tag_var,dim=1)
100 | 
101 |             t_r1_k = torch.unsqueeze(feats[feat_index], 0)
102 |             forward_var_new = torch.unsqueeze(viterbivars_t,0) + t_r1_k
103 | 
104 |             forward_var_list.append(forward_var_new)
105 |             backpointers.append(bptrs_t.tolist())
106 | 
107 |         # Transition to STOP_TAG
108 |         terminal_var = forward_var_list[-1] + self.transitions[self.tag_to_ix[STOP_TAG]]
109 |         best_tag_id = torch.argmax(terminal_var).tolist()
110 |         path_score = terminal_var[0][best_tag_id]
111 | 
112 |         # Follow the back pointers to decode the best path.
113 |         best_path = [best_tag_id]
114 |         for bptrs_t in reversed(backpointers):
115 |             best_tag_id = bptrs_t[best_tag_id]
116 |             best_path.append(best_tag_id)
117 |         # Pop off the start tag (we dont want to return that to the caller)
118 |         start = best_path.pop()
119 |         assert start == self.tag_to_ix[START_TAG]  # Sanity check
120 |         best_path.reverse()
121 |         return path_score, best_path
122 | 
123 | 
124 |     def neg_log_likelihood(self, sentence, tags):
125 |         feats = self._get_lstm_features(sentence)
126 |         forward_score = self._forward_alg(feats)
127 |         gold_score = self._score_sentence(feats, tags)
128 |         return forward_score - gold_score
129 | 
130 |     def forward(self, sentence):  # dont confuse this with _forward_alg above.
131 |         # Get the emission scores from the BiLSTM
132 |         lstm_feats = self._get_lstm_features(sentence)
133 | 
134 |         # Find the best path, given the features.
135 |         score, tag_seq = self._viterbi_decode(lstm_feats)
136 |         return score, tag_seq
137 | 
138 | 
139 | if __name__== '__main__':
140 |     START_TAG = "<START>"
141 |     STOP_TAG = "<STOP>"
142 |     EMBEDDING_DIM = 300
143 |     HIDDEN_DIM = 256
144 | 
145 |     # Make up some training data
146 |     training_data = [(
147 |         "the wall street journal reported today that apple corporation made money".split(),
148 |         "B I I I O O O B I O O".split()
149 |     ), (
150 |         "georgia tech is a university in georgia".split(),
151 |         "B I O O O O B".split()
152 |     )]
153 | 
154 |     word_to_ix = {}
155 |     for sentence, tags in training_data:
156 |         for word in sentence:
157 |             if word not in word_to_ix:
158 |                 word_to_ix[word] = len(word_to_ix)
159 | 
160 |     tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}
161 | 
162 |     model = BiLSTM_CRF(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)
163 |     optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)
164 | 
165 |     # Check predictions before training
166 |     with torch.no_grad():
167 |         precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
168 |         precheck_tags = torch.tensor([tag_to_ix[t] for t in training_data[0][1]], dtype=torch.long)
169 |         print(model(precheck_sent))
170 | 
171 |     # Make sure prepare_sequence from earlier in the LSTM section is loaded
172 |     for epoch in range(
173 |             300):  # again, normally you would NOT do 300 epochs, it is toy data
174 |         for sentence, tags in training_data:
175 |             # Step 1. Remember that Pytorch accumulates gradients.
176 |             # We need to clear them out before each instance
177 |             model.zero_grad()
178 | 
179 |             # Step 2. Get our inputs ready for the network, that is,
180 |             # turn them into Tensors of word indices.
181 |             sentence_in = prepare_sequence(sentence, word_to_ix)
182 |             targets = torch.tensor([tag_to_ix[t] for t in tags], dtype=torch.long)
183 | 
184 |             # Step 3. Run our forward pass.
185 |             loss = model.neg_log_likelihood(sentence_in, targets)
186 | 
187 |             # Step 4. Compute the loss, gradients, and update the parameters by
188 |             # calling optimizer.step()
189 |             loss.backward()
190 |             optimizer.step()
191 | 
192 |     # Check predictions after training
193 |     with torch.no_grad():
194 |         precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
195 |         print(model(precheck_sent))
196 |         # We got it!


--------------------------------------------------------------------------------
/LSTM_CRF_faster_parallel.py:
--------------------------------------------------------------------------------
  1 | import time
  2 | import torch
  3 | import torch.nn as nn
  4 | import torch.optim as optim
  5 | 
  6 | START_TAG = "<START>"
  7 | STOP_TAG = "<STOP>"
  8 | 
  9 | # torch.manual_seed(1)
 10 | 
 11 | def argmax(vec):
 12 |     # return the argmax as a python int
 13 |     _, idx = torch.max(vec, 1)
 14 |     return idx.item()
 15 | 
 16 | 
 17 | def prepare_sequence(seq, to_ix):
 18 |     idxs = [to_ix[w] for w in seq]
 19 |     return torch.tensor(idxs, dtype=torch.long)
 20 | 
 21 | def prepare_sequence_batch(data ,word_to_ix, tag_to_ix):
 22 |     seqs = [i[0] for i in data]
 23 |     tags = [i[1] for i in data]
 24 |     max_len = max([len(seq) for seq in seqs])
 25 |     seqs_pad=[]
 26 |     tags_pad=[]
 27 |     for seq,tag in zip(seqs, tags):
 28 |         seq_pad = seq + ['<PAD>'] * (max_len-len(seq))
 29 |         tag_pad = tag + ['<PAD>'] * (max_len-len(tag))
 30 |         seqs_pad.append(seq_pad)
 31 |         tags_pad.append(tag_pad)
 32 |     idxs_pad = torch.tensor([[word_to_ix[w] for w in seq] for seq in seqs_pad], dtype=torch.long)
 33 |     tags_pad = torch.tensor([[tag_to_ix[t] for t in tag] for tag in tags_pad], dtype=torch.long)
 34 |     return idxs_pad, tags_pad
 35 | 
 36 | 
 37 | # Compute log sum exp in a numerically stable way for the forward algorithm
 38 | def log_sum_exp(vec):
 39 |     max_score = vec[0, argmax(vec)]
 40 |     max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
 41 |     return max_score + \
 42 |            torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))
 43 | 
 44 | 
 45 | def log_add(args):
 46 |     return torch.log(torch.sum(torch.exp(args), axis=0))
 47 | 
 48 | 
 49 | 
 50 | class BiLSTM_CRF_MODIFY_PARALLEL(nn.Module):
 51 | 
 52 |     def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):
 53 |         super(BiLSTM_CRF_MODIFY_PARALLEL, self).__init__()
 54 |         self.embedding_dim = embedding_dim
 55 |         self.hidden_dim = hidden_dim
 56 |         self.vocab_size = vocab_size
 57 |         self.tag_to_ix = tag_to_ix
 58 |         self.tagset_size = len(tag_to_ix)
 59 | 
 60 |         self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
 61 |         self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2,
 62 |                             num_layers=1, bidirectional=True, batch_first=True)
 63 | 
 64 |         # Maps the output of the LSTM into tag space.
 65 |         self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
 66 | 
 67 |         # Matrix of transition parameters.  Entry i,j is the score of
 68 |         # transitioning *to* i *from* j.
 69 |         self.transitions = nn.Parameter(
 70 |             torch.randn(self.tagset_size, self.tagset_size))
 71 | 
 72 |         # These two statements enforce the constraint that we never transfer
 73 |         # to the start tag and we never transfer from the stop tag
 74 | 
 75 |         self.transitions.data[tag_to_ix[START_TAG], :] = -10000
 76 |         self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000
 77 |         self.hidden = self.init_hidden()
 78 | 
 79 |     def init_hidden(self):
 80 |         return (torch.randn(2, 1, self.hidden_dim // 2),
 81 |                 torch.randn(2, 1, self.hidden_dim // 2))
 82 | 
 83 |     def _forward_alg(self, feats):
 84 |         begin = time.time()
 85 |         # Do the forward algorithm to compute the partition function
 86 |         init_alphas = torch.full((1, self.tagset_size), -10000.).to('cuda')
 87 |         # START_TAG has all of the score.
 88 |         init_alphas[0][self.tag_to_ix[START_TAG]] = 0.
 89 | 
 90 |         # Wrap in a variable so that we will get automatic backprop
 91 |         forward_var = init_alphas
 92 |         # print('time consuming of crf_partion_function_prepare:%f' % (time.time() - begin))
 93 |         begin = time.time()
 94 |         # Iterate through the sentence
 95 |         for feat in feats:
 96 |             alphas_t = []  # The forward tensors at this timestep
 97 |             for next_tag in range(self.tagset_size):
 98 |                 # broadcast the emission score: it is the same regardless of
 99 |                 # the previous tag
100 |                 emit_score = feat[next_tag].view(
101 |                     1, -1).expand(1, self.tagset_size)
102 |                 # the ith entry of trans_score is the score of transitioning to
103 |                 # next_tag from i
104 |                 trans_score = self.transitions[next_tag].view(1, -1)
105 |                 # The ith entry of next_tag_var is the value for the
106 |                 # edge (i -> next_tag) before we do log-sum-exp
107 |                 next_tag_var = (forward_var + trans_score + emit_score)
108 |                 # The forward variable for this tag is log-sum-exp of all the
109 |                 # scores.
110 |                 alphas_t.append(log_sum_exp(next_tag_var).view(1))
111 |             forward_var = torch.cat(alphas_t).view(1, -1)
112 |         # print('time consuming of crf_partion_function1:%f' % (time.time() - begin))
113 |         terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
114 |         alpha = log_sum_exp(terminal_var)
115 |         # print('time consuming of crf_partion_function2:%f' %(time.time()-begin))
116 |         return alpha
117 | 
118 |     def _forward_alg_new(self, feats):
119 |         # Do the forward algorithm to compute the partition function
120 |         init_alphas = torch.full([self.tagset_size], -10000.).to('cuda')
121 |         # START_TAG has all of the score.
122 |         init_alphas[self.tag_to_ix[START_TAG]] = 0.
123 | 
124 |         # Wrap in a variable so that we will get automatic backprop
125 |         # Iterate through the sentence
126 |         forward_var_list = []
127 |         forward_var_list.append(init_alphas)
128 |         for feat_index in range(feats.shape[0]):  # -1
129 |             gamar_r_l = torch.stack([forward_var_list[feat_index]] * feats.shape[1])
130 |             # gamar_r_l = torch.transpose(gamar_r_l,0,1)
131 |             t_r1_k = torch.unsqueeze(feats[feat_index], 0).transpose(0, 1)  # +1
132 |             aa = gamar_r_l + t_r1_k + self.transitions
133 |             # forward_var_list.append(log_add(aa))
134 |             forward_var_list.append(torch.logsumexp(aa, dim=1))
135 |         terminal_var = forward_var_list[-1] + self.transitions[self.tag_to_ix[STOP_TAG]]
136 |         terminal_var = torch.unsqueeze(terminal_var, 0)
137 |         alpha = torch.logsumexp(terminal_var, dim=1)[0]
138 |         return alpha
139 | 
140 |     def _forward_alg_new_parallel(self, feats):
141 |         # Do the forward algorithm to compute the partition function
142 |         init_alphas = torch.full([feats.shape[0], self.tagset_size], -10000.)#.to('cuda')
143 |         # START_TAG has all of the score.
144 |         init_alphas[:, self.tag_to_ix[START_TAG]] = 0.
145 | 
146 |         # Wrap in a variable so that we will get automatic backprop
147 |         # Iterate through the sentence
148 |         forward_var_list = []
149 |         forward_var_list.append(init_alphas)
150 |         for feat_index in range(feats.shape[1]):  # -1
151 |             gamar_r_l = torch.stack([forward_var_list[feat_index]] * feats.shape[2]).transpose(0, 1)
152 |             # gamar_r_l = torch.transpose(gamar_r_l,0,1)
153 |             t_r1_k = torch.unsqueeze(feats[:, feat_index, :], 1).transpose(1, 2)  # +1
154 |             # t_r1_k = feats[:,feat_index,:].repeat(feats.shape[0],1,1).transpose(1, 2)
155 |             aa = gamar_r_l + t_r1_k + torch.unsqueeze(self.transitions, 0)
156 |             # forward_var_list.append(log_add(aa))
157 |             forward_var_list.append(torch.logsumexp(aa, dim=2))
158 |         terminal_var = forward_var_list[-1] + self.transitions[self.tag_to_ix[STOP_TAG]].repeat([feats.shape[0], 1])
159 |         # terminal_var = torch.unsqueeze(terminal_var, 0)
160 |         alpha = torch.logsumexp(terminal_var, dim=1)
161 |         return alpha
162 | 
163 | 
164 |     def _get_lstm_features(self, sentence):
165 |         self.hidden = self.init_hidden()
166 |         embeds = self.word_embeds(sentence).unsqueeze(dim=0)
167 |         #embeds = self.word_embeds(sentence).view(len(sentence), 1, -1).transpose(0,1)
168 |         lstm_out, self.hidden = self.lstm(embeds)
169 |         #lstm_out = lstm_out.view(embeds.shape[1], self.hidden_dim)
170 |         lstm_out = lstm_out.squeeze()
171 |         lstm_feats = self.hidden2tag(lstm_out)
172 |         return lstm_feats
173 | 
174 |     def _get_lstm_features_parallel(self, sentence):
175 |         self.hidden = self.init_hidden()
176 |         embeds = self.word_embeds(sentence)
177 |         lstm_out, self.hidden = self.lstm(embeds)
178 |         lstm_feats = self.hidden2tag(lstm_out)
179 |         return lstm_feats
180 | 
181 |     def _score_sentence(self, feats, tags):
182 |         # Gives the score of a provided tag sequence
183 |         score = torch.zeros(1)
184 |         # score = autograd.Variable(torch.Tensor([0])).to('cuda')
185 |         tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags.view(-1)])
186 | 
187 |         # if len(tags)<2:
188 |         #     print(tags)
189 |         #     sys.exit(0)
190 |         for i, feat in enumerate(feats):
191 |             score = score + \
192 |                     self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
193 |         score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
194 |         return score
195 | 
196 |     def _score_sentence_parallel(self, feats, tags):
197 |         # Gives the score of provided tag sequences
198 |         #feats = feats.transpose(0,1)
199 | 
200 |         score = torch.zeros(tags.shape[0])#.to('cuda')
201 |         tags = torch.cat([torch.full([tags.shape[0],1],self.tag_to_ix[START_TAG], dtype=torch.long),tags],dim=1)
202 |         for i in range(feats.shape[1]):
203 |             feat=feats[:,i,:]
204 |             score = score + \
205 |                     self.transitions[tags[:,i + 1], tags[:,i]] + feat[range(feat.shape[0]),tags[:,i + 1]]
206 |         score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[:,-1]]
207 |         return score
208 | 
209 | 
210 | 
211 |     def _viterbi_decode(self, feats):
212 |         backpointers = []
213 | 
214 |         # Initialize the viterbi variables in log space
215 |         init_vvars = torch.full((1, self.tagset_size), -10000.)
216 |         init_vvars[0][self.tag_to_ix[START_TAG]] = 0
217 | 
218 |         # forward_var at step i holds the viterbi variables for step i-1
219 |         forward_var = init_vvars
220 | 
221 |         for feat in feats:
222 |             bptrs_t = []  # holds the backpointers for this step
223 |             viterbivars_t = []  # holds the viterbi variables for this step
224 | 
225 |             for next_tag in range(self.tagset_size):
226 |                 # next_tag_var[i] holds the viterbi variable for tag i at the
227 |                 # previous step, plus the score of transitioning
228 |                 # from tag i to next_tag.
229 |                 # We don't include the emission scores here because the max
230 |                 # does not depend on them (we add them in below)
231 |                 next_tag_var = forward_var.to('cuda') + self.transitions[next_tag]
232 |                 best_tag_id = argmax(next_tag_var)
233 |                 bptrs_t.append(best_tag_id)
234 |                 viterbivars_t.append(next_tag_var[0][best_tag_id].view(1))
235 |             # Now add in the emission scores, and assign forward_var to the set
236 |             # of viterbi variables we just computed
237 |             forward_var = (torch.cat(viterbivars_t) + feat).view(1, -1)
238 |             backpointers.append(bptrs_t)
239 | 
240 |         # Transition to STOP_TAG
241 |         terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
242 |         best_tag_id = argmax(terminal_var)
243 |         path_score = terminal_var[0][best_tag_id]
244 | 
245 |         # Follow the back pointers to decode the best path.
246 |         best_path = [best_tag_id]
247 |         for bptrs_t in reversed(backpointers):
248 |             best_tag_id = bptrs_t[best_tag_id]
249 |             best_path.append(best_tag_id)
250 |         # Pop off the start tag (we dont want to return that to the caller)
251 |         start = best_path.pop()
252 |         assert start == self.tag_to_ix[START_TAG]  # Sanity check
253 |         best_path.reverse()
254 |         return path_score, best_path
255 | 
256 |     def _viterbi_decode_new(self, feats):
257 |         backpointers = []
258 | 
259 |         # Initialize the viterbi variables in log space
260 |         init_vvars = torch.full((1, self.tagset_size), -10000.)#.to('cuda')
261 |         init_vvars[0][self.tag_to_ix[START_TAG]] = 0
262 | 
263 |         # forward_var at step i holds the viterbi variables for step i-1
264 |         forward_var_list = []
265 |         forward_var_list.append(init_vvars)
266 | 
267 |         for feat_index in range(feats.shape[0]):
268 |             gamar_r_l = torch.stack([forward_var_list[feat_index]] * feats.shape[1])
269 |             gamar_r_l = torch.squeeze(gamar_r_l)
270 |             next_tag_var = gamar_r_l + self.transitions
271 |             # bptrs_t=torch.argmax(next_tag_var,dim=0)
272 |             viterbivars_t, bptrs_t = torch.max(next_tag_var, dim=1)
273 | 
274 |             t_r1_k = torch.unsqueeze(feats[feat_index], 0)
275 |             forward_var_new = torch.unsqueeze(viterbivars_t, 0) + t_r1_k
276 | 
277 |             forward_var_list.append(forward_var_new)
278 |             backpointers.append(bptrs_t.tolist())
279 | 
280 |         # Transition to STOP_TAG
281 |         terminal_var = forward_var_list[-1] + self.transitions[self.tag_to_ix[STOP_TAG]]
282 |         best_tag_id = torch.argmax(terminal_var).tolist()
283 |         path_score = terminal_var[0][best_tag_id]
284 | 
285 |         # Follow the back pointers to decode the best path.
286 |         best_path = [best_tag_id]
287 |         for bptrs_t in reversed(backpointers):
288 |             best_tag_id = bptrs_t[best_tag_id]
289 |             best_path.append(best_tag_id)
290 |         # Pop off the start tag (we dont want to return that to the caller)
291 |         start = best_path.pop()
292 |         assert start == self.tag_to_ix[START_TAG]  # Sanity check
293 |         best_path.reverse()
294 |         return path_score, best_path
295 | 
296 |     def neg_log_likelihood(self, sentence, tags):
297 |         feats = self._get_lstm_features(sentence)
298 |         forward_score = self._forward_alg_new(feats)
299 |         gold_score = self._score_sentence(feats, tags)[0]
300 |         return forward_score - gold_score
301 | 
302 |     def neg_log_likelihood_parallel(self, sentences, tags):
303 |         feats = self._get_lstm_features_parallel(sentences)
304 |         forward_score = self._forward_alg_new_parallel(feats)
305 |         gold_score = self._score_sentence_parallel(feats, tags)
306 |         return torch.sum(forward_score - gold_score)
307 | 
308 |     def forward(self, sentence):  # dont confuse this with _forward_alg above.
309 |         # Get the emission scores from the BiLSTM
310 |         lstm_feats = self._get_lstm_features(sentence)
311 | 
312 |         # Find the best path, given the features.
313 |         score, tag_seq = self._viterbi_decode_new(lstm_feats)
314 |         return score, tag_seq
315 | 
316 | if __name__ == '__main__':
317 |     START_TAG = "<START>"
318 |     STOP_TAG = "<STOP>"
319 |     PAD_TAG = "<PAD>"
320 |     EMBEDDING_DIM = 300
321 |     HIDDEN_DIM = 256
322 | 
323 |     # Make up some training data
324 |     training_data = [(
325 |         "the wall street journal reported today that apple corporation made money".split(),
326 |         "B I I I O O O B I O O".split()
327 |     ), (
328 |         "georgia tech is a university in georgia".split(),
329 |         "B I O O O O B".split()
330 |     )]
331 | 
332 |     word_to_ix = {}
333 |     word_to_ix['<PAD>'] = 0
334 |     for sentence, tags in training_data:
335 |         for word in sentence:
336 |             if word not in word_to_ix:
337 |                 word_to_ix[word] = len(word_to_ix)
338 | 
339 |     tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4, PAD_TAG: 5}
340 | 
341 |     model = BiLSTM_CRF_MODIFY_PARALLEL(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)
342 |     optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)
343 | 
344 |     # Check predictions before training
345 |     with torch.no_grad():
346 |         precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
347 |         precheck_tags = torch.tensor([tag_to_ix[t] for t in training_data[0][1]], dtype=torch.long)
348 |         print(model(precheck_sent))
349 | 
350 |     # Make sure prepare_sequence from earlier in the LSTM section is loaded
351 |     for epoch in range(
352 |             300):  # again, normally you would NOT do 300 epochs, it is toy data
353 |         # Step 1. Remember that Pytorch accumulates gradients.
354 |         # We need to clear them out before each instance
355 |         model.zero_grad()
356 |         # Step 2. Get our batch inputs ready for the network, that is,
357 |         # turn them into Tensors of word indices.
358 |         # If training_data can't be included in one batch, you need to sample them to build a batch
359 |         sentence_in_pad, targets_pad = prepare_sequence_batch(training_data, word_to_ix, tag_to_ix)
360 |         # Step 3. Run our forward pass.
361 |         loss = model.neg_log_likelihood_parallel(sentence_in_pad, targets_pad)
362 |         # Step 4. Compute the loss, gradients, and update the parameters by
363 |         # calling optimizer.step()
364 |         loss.backward()
365 |         optimizer.step()
366 | 
367 |     # Check predictions after training
368 |     with torch.no_grad():
369 |         precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
370 |         print(model(precheck_sent))
371 |         # We got it!
372 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # LSTM-CRF-pytorch-faster
 2 | 
 3 | This is a more than 1000X faster LSTM-CRF implementation modified from the slower version in the Pytorch official tutorial   (URL:https://pytorch.org/tutorials/beginner/nlp/advanced_tutorial.html). 
 4 | 
 5 | I have modified the dynamic planning part, including Viterbi decoding and partition function calculation. In the experiment, it has achieved a speed increase of more than 50 times compared to the original version. Furthermore, the original version can only input one sample at a time when it is trained. 
 6 | 
 7 | In the most recent updated module 'LSTM_CRF_faster_parallel.py', I modified the model to support parallel computing for batch, so that the training time was greatly improved again. When the batchsize is large, parallel computing can bring you hundreds of times faster.
 8 | 
 9 | The code defaults to training word embedding from scratch. If you need to use pre-training word embedding, or take another model's outputs as it's inputs, you need to make some changes to the code, directly take the embedding as inpus of method '_get_lstm_features_parallel'.
10 | 
11 | Moreover, the code runs on the CPU by default, and if you need to experience faster speed, you need to manually deploy the model to the GPU.
12 |  
13 | (In previous version, parallel version 'LSTM_CRF_faster_parallel.py' is not work, now the bug is rectified. For more questions please send email to M201570013@outlook.com)
14 | 
15 | 
16 | 


--------------------------------------------------------------------------------