def eval_one_batch(self, batch): enc_batch, enc_lens, enc_pos, enc_padding_mask, enc_batch_extend_vocab, extra_zeros, c_t, coverage = \ get_input_from_batch(batch, use_cuda) dec_batch, dec_lens, dec_pos, dec_padding_mask, max_dec_len, tgt_batch = \ get_output_from_batch(batch, use_cuda) enc_out, enc_fea, enc_h = self.model.encoder(enc_batch, enc_lens) s_t = self.model.reduce_state(enc_h) step_losses = [] for di in range(min(max_dec_len, config.max_dec_steps)): y_t = dec_batch[:, di] # Teacher forcing final_dist, s_t, c_t, attn_dist, p_gen, next_coverage = self.model.decoder( y_t, s_t, enc_out, enc_fea, enc_padding_mask, c_t, extra_zeros, enc_batch_extend_vocab, coverage, di) tgt = tgt_batch[:, di] gold_probs = torch.gather(final_dist, 1, tgt.unsqueeze(1)).squeeze() step_loss = -torch.log(gold_probs + config.eps) if config.is_coverage: step_coverage_loss = torch.sum(torch.min(attn_dist, coverage), 1) step_loss = step_loss + config.cov_loss_wt * step_coverage_loss coverage = next_coverage step_mask = dec_padding_mask[:, di] step_loss = step_loss * step_mask step_losses.append(step_loss) sum_step_losses = torch.sum(torch.stack(step_losses, 1), 1) batch_avg_loss = sum_step_losses / dec_lens loss = torch.mean(batch_avg_loss) return loss.item()
def train_one_batch(self, batch): enc_batch, enc_lens, enc_pos, enc_padding_mask, enc_batch_extend_vocab, \ extra_zeros, c_t, coverage = get_input_from_batch(batch, use_cuda) dec_batch, dec_lens, dec_pos, dec_padding_mask, max_dec_len, tgt_batch = \ get_output_from_batch(batch, use_cuda) self.optimizer.zero_grad() if not config.tran: enc_out, enc_fea, enc_h = self.model.encoder(enc_batch, enc_lens) else: enc_out, enc_fea, enc_h = self.model.encoder(enc_batch, enc_pos) s_t = self.model.reduce_state(enc_h) step_losses, cove_losses = [], [] for di in range(min(max_dec_len, config.max_dec_steps)): y_t = dec_batch[:, di] # Teacher forcing final_dist, s_t, c_t, attn_dist, p_gen, next_coverage = \ self.model.decoder(y_t, s_t, enc_out, enc_fea, enc_padding_mask, c_t, extra_zeros, enc_batch_extend_vocab, coverage, di) tgt = tgt_batch[:, di] step_mask = dec_padding_mask[:, di] gold_probs = torch.gather(final_dist, 1, tgt.unsqueeze(1)).squeeze() step_loss = -torch.log(gold_probs + config.eps) if config.is_coverage: step_coverage_loss = torch.sum(torch.min(attn_dist, coverage), 1) step_loss = step_loss + config.cov_loss_wt * step_coverage_loss cove_losses.append(step_coverage_loss * step_mask) coverage = next_coverage step_loss = step_loss * step_mask step_losses.append(step_loss) sum_losses = torch.sum(torch.stack(step_losses, 1), 1) batch_avg_loss = sum_losses / dec_lens loss = torch.mean(batch_avg_loss) loss.backward() clip_grad_norm_(self.model.encoder.parameters(), config.max_grad_norm) clip_grad_norm_(self.model.decoder.parameters(), config.max_grad_norm) clip_grad_norm_(self.model.reduce_state.parameters(), config.max_grad_norm) self.optimizer.step() if config.is_coverage: cove_losses = torch.sum(torch.stack(cove_losses, 1), 1) batch_cove_loss = cove_losses / dec_lens batch_cove_loss = torch.mean(batch_cove_loss) return loss.item(), batch_cove_loss.item() return loss.item(), 0.
def train_one_batch(self, batch): enc_batch, enc_lens, enc_pos, enc_padding_mask, enc_batch_extend_vocab, \ extra_zeros, c_t, coverage = get_input_from_batch(batch, use_cuda, transformer=True) dec_batch, dec_lens, dec_pos, dec_padding_mask, max_dec_len, tgt_batch = \ get_output_from_batch(batch, use_cuda, transformer=True) self.optimizer.zero_grad() pred = self.model(enc_batch, enc_pos, dec_batch, dec_pos) gold_probs = torch.gather(pred, -1, tgt_batch.unsqueeze(-1)).squeeze() batch_loss = -torch.log(gold_probs + config.eps) batch_loss = batch_loss * dec_padding_mask sum_losses = torch.sum(batch_loss, 1) batch_avg_loss = sum_losses / dec_lens loss = torch.mean(batch_avg_loss) loss.backward() # update parameters self.optimizer.step_and_update_lr() return loss.item(), 0.
def train_one_batch(self, batch): enc_batch, enc_lens, enc_pos, enc_padding_mask, enc_batch_extend_vocab, \ extra_zeros, c_t, coverage = get_input_from_batch(batch, use_cuda) dec_batch, dec_lens, dec_pos, dec_padding_mask, max_dec_len, tgt_batch = \ get_output_from_batch(batch, use_cuda) self.optimizer.zero_grad() enc_outs = [] enc_feas = [] enc_hs = [] if not config.tran: for i, encoder in enumerate(self.model.encoders): enc_out, enc_fea, enc_h = encoder(enc_batch[i], enc_lens[i]) enc_outs.append(enc_out) enc_feas.append(enc_fea) enc_hs.append(enc_h) # else: # enc_out, enc_fea, enc_h = self.model.encoder(enc_batch, enc_pos) s_t = self.model.reduce_state(enc_h) step_losses, cove_losses = [], [] for di in range(min(max_dec_len, config.max_dec_steps)): y_t = dec_batch[:, di] # Teacher forcing # modify the original frame for two encoders. final_dist_0, s_t_0, c_t_0, attn_dist_0, p_gen, next_coverage = \ self.model.decoder(y_t, s_t, enc_outs[0], enc_feas[0], enc_padding_mask[0], c_t, extra_zeros, enc_batch_extend_vocab[0], coverage, di) final_dist_1, s_t_1, c_t_1, attn_dist, p_gen, next_coverage = \ self.model.decoder(y_t, s_t, enc_outs[1], enc_feas[1], enc_padding_mask[1], c_t, extra_zeros, enc_batch_extend_vocab[1], coverage, di) y_t_emb = self.model.decoder.tgt_word_emb(y_t) encoders_att = self.model.encoders_att(enc_hs, y_t_emb) final_dist = torch.stack((final_dist_0, final_dist_1), dim=1) final_dist = torch.bmm(encoders_att, final_dist).squeeze() encoders_att_ = encoders_att.transpose(0, 1).contiguous() # 1 x b x 2 h = s_t_0[0] * encoders_att_[:, :, :1] + s_t_1[ 0] * encoders_att_[:, :, 1:] c = s_t_0[1] * encoders_att_[:, :, :1] + s_t_1[ 1] * encoders_att_[:, :, 1:] s_t = (h, c) tgt = tgt_batch[:, di] step_mask = dec_padding_mask[:, di] gold_probs = torch.gather(final_dist, 1, tgt.unsqueeze(1)).squeeze() step_loss = -torch.log(gold_probs + config.eps) if config.is_coverage: step_coverage_loss = torch.sum(torch.min(attn_dist, coverage), 1) step_loss = step_loss + config.cov_loss_wt * step_coverage_loss cove_losses.append(step_coverage_loss * step_mask) coverage = next_coverage step_loss = step_loss * step_mask step_losses.append(step_loss) sum_losses = torch.sum(torch.stack(step_losses, 1), 1) batch_avg_loss = sum_losses / dec_lens loss = torch.mean(batch_avg_loss) loss.backward() clip_grad_norm_(self.model.encoders.parameters(), config.max_grad_norm) clip_grad_norm_(self.model.decoder.parameters(), config.max_grad_norm) clip_grad_norm_(self.model.reduce_state.parameters(), config.max_grad_norm) self.optimizer.step() if config.is_coverage: cove_losses = torch.sum(torch.stack(cove_losses, 1), 1) batch_cove_loss = cove_losses / dec_lens batch_cove_loss = torch.mean(batch_cove_loss) return loss.item(), batch_cove_loss.item() return loss.item(), 0.
def beam_search(self, batch): # single example repeated across the batch enc_batch, enc_lens, enc_pos, enc_padding_mask, enc_batch_extend_vocab, extra_zeros, c_t, coverage = \ get_input_from_batch(batch, use_cuda) enc_out, enc_fea, enc_h = self.model.encoder(enc_batch, enc_lens) s_t = self.model.reduce_state(enc_h) dec_h, dec_c = s_t # b x hidden_dim dec_h = dec_h.squeeze() dec_c = dec_c.squeeze() # decoder batch preparation, it has beam_size example initially everything is repeated beams = [ Beam(tokens=[self.vocab.word2id(config.BOS_TOKEN)], log_probs=[0.0], state=(dec_h[0], dec_c[0]), context=c_t[0], coverage=(coverage[0] if config.is_coverage else None)) for _ in range(config.beam_size) ] steps = 0 results = [] while steps < config.max_dec_steps and len(results) < config.beam_size: latest_tokens = [h.latest_token for h in beams] latest_tokens = [t if t < self.vocab.size() else self.vocab.word2id(config.UNK_TOKEN) \ for t in latest_tokens] y_t = Variable(torch.LongTensor(latest_tokens)) if use_cuda: y_t = y_t.cuda() all_state_h = [h.state[0] for h in beams] all_state_c = [h.state[1] for h in beams] all_context = [h.context for h in beams] s_t = (torch.stack(all_state_h, 0).unsqueeze(0), torch.stack(all_state_c, 0).unsqueeze(0)) c_t = torch.stack(all_context, 0) coverage_t = None if config.is_coverage: all_coverage = [h.coverage for h in beams] coverage_t = torch.stack(all_coverage, 0) final_dist, s_t, c_t, attn_dist, p_gen, coverage_t = self.model.decoder( y_t, s_t, enc_out, enc_fea, enc_padding_mask, c_t, extra_zeros, enc_batch_extend_vocab, coverage_t, steps) log_probs = torch.log(final_dist) topk_log_probs, topk_ids = torch.topk(log_probs, config.beam_size * 2) dec_h, dec_c = s_t dec_h = dec_h.squeeze() dec_c = dec_c.squeeze() all_beams = [] # On the first step, we only had one original hypothesis (the initial hypothesis). On subsequent steps, all original hypotheses are distinct. num_orig_beams = 1 if steps == 0 else len(beams) for i in range(num_orig_beams): h = beams[i] state_i = (dec_h[i], dec_c[i]) context_i = c_t[i] coverage_i = (coverage[i] if config.is_coverage else None) for j in range(config.beam_size * 2): # for each of the top 2*beam_size hyps: new_beam = h.extend(token=topk_ids[i, j].item(), log_prob=topk_log_probs[i, j].item(), state=state_i, context=context_i, coverage=coverage_i) all_beams.append(new_beam) beams = [] for h in self.sort_beams(all_beams): if h.latest_token == self.vocab.word2id(config.EOS_TOKEN): if steps >= config.min_dec_steps: results.append(h) else: beams.append(h) if len(beams) == config.beam_size or len( results) == config.beam_size: break steps += 1 if len(results) == 0: results = beams beams_sorted = self.sort_beams(results) return beams_sorted[0]