seq2seq BeamSearch 구현

iml1111 · iml1111 · commit de856f2f2ab8 · 2021-01-17T00:35:48.000+09:00
diff --git a/src/11_seq2seq/modules/search.py b/src/11_seq2seq/modules/search.py
@@ -0,0 +1,217 @@
+from operator import itemgetter
+
+import torch
+import torch.nn as nn
+
+import modules.data_loader as data_loader
+
+LENGTH_PENALTY = .2
+MIN_LENGTH = 5
+
+
+class SingleBeamSearchBoard():
+
+    def __init__(
+        self,
+        device,
+        prev_status_config,
+        beam_size=5,
+        max_length=255,
+    ):
+        self.beam_size = beam_size
+        self.max_length = max_length
+
+        self.device = device
+        # 각 타임 스텝의 Word Index(즉, 최종 예측 단어들) * beam_size
+        # 처음에는 모두 <BOS>므로 초기화
+        self.word_indice = [torch.LongTensor(beam_size).zero_().to(self.device) + data_loader.BOS]
+        # 각 타임 스텝의 Word들이 선정된 Beam Index
+        # 처음에는 아무것도 선정되지 않았기에 -1로 초기화
+        self.beam_indice = [torch.LongTensor(beam_size).zero_().to(self.device) - 1]
+        # 각 Beam들의 누적 확률 값
+        # 처음에는 [0, -inf, -inf, ...]로 초기화
+        self.cumulative_probs = [torch.FloatTensor([.0] + [-float('inf')] * (beam_size - 1)).to(self.device)]
+        # 각 빔이 현재 EOS에 도달했는지 여부
+        # 1 if it is done else 0
+        self.masks = [torch.BoolTensor(beam_size).zero_().to(self.device)]
+        # We don't need to remember every time-step of hidden states:
+        #       prev_hidden, prev_cell, prev_h_t_tilde
+        # What we need is remember just last one.
+
+        '''
+        각 빔의 이전 hidden, cell, h_tilde를 저장해두는 공간
+        항상 마지막 타임스텝만 보관하면 됨
+
+        단 처음에는 그냥 넘겨받은 hidden, cell, h_tilde를 beam_size만큼 늘려줌
+        h_tilde의 경우, 처음에 None이므로 예외처리
+        '''
+        self.prev_status = {}
+        self.batch_dims = {}
+        for prev_status_name, each_config in prev_status_config.items():
+            init_status = each_config['init_status']
+            batch_dim_index = each_config['batch_dim_index']
+            if init_status is not None:
+                self.prev_status[prev_status_name] = torch.cat([init_status] * beam_size,
+                                                               dim=batch_dim_index)
+            else:
+                self.prev_status[prev_status_name] = None
+            self.batch_dims[prev_status_name] = batch_dim_index
+
+        self.current_time_step = 0
+        self.done_cnt = 0
+
+    def get_length_penalty(
+        self,
+        length,
+        alpha=LENGTH_PENALTY,
+        min_length=MIN_LENGTH,
+    ):
+        # Calculate length-penalty,
+        # because shorter sentence usually have bigger probability.
+        # In fact, we represent this as log-probability, which is negative value.
+        # Thus, we need to multiply bigger penalty for shorter one.
+        p = ((min_length + 1) / (min_length + length))**alpha
+
+        return p
+
+    def is_done(self):
+        # Return 1, if we had EOS more than 'beam_size'-times.
+        if self.done_cnt >= self.beam_size:
+            return 1
+        return 0
+
+    def get_batch(self):
+        '''
+        현재 빔에서 가장 마지막 스텝의 워드 인덱스들을 가져옴
+        처음에는, 당연히 모두 BOS 일것임
+        그 후로는 이전에 예측했던 TopK의 단어들을 주게 될것임
+        '''
+        y_hat = self.word_indice[-1].unsqueeze(-1)
+        # |y_hat| = (beam_size, 1)
+        # if model != transformer:
+        #     |hidden| = |cell| = (n_layers, beam_size, hidden_size)
+        #     |h_t_tilde| = (beam_size, 1, hidden_size) or None
+        # else:
+        #     |prev_state_i| = (beam_size, length, hidden_size),
+        #     where i is an index of layer.
+        return y_hat, self.prev_status
+
+    #@profile
+    def collect_result(self, y_hat, prev_status):
+        '''
+        y_hat: 현재 타입스텝의 각 beam마다 예측한 단어
+        pre_status: 현재 타입스텝에서 함께 나왔던 hidden, cell, h_tilde
+        넣을때, beam 채로 넣었으므로 그대로 다시 나오게 됨
+        '''
+
+        # |y_hat| = (beam_size, 1, output_size)
+        # |hidden| = |cell| = (n_layers, beam_size, hidden_size)
+        # |h_t_tilde| = (beam_size, 1, hidden_size)
+        output_size = y_hat.size(-1)
+
+        self.current_time_step += 1
+
+        # 누적 확률 값을 계산함
+        # (beam_size) --> (beam_size, 1, 1) --> (beam_size, 1, output_size)
+        # 이미 예측이 끝난 경우, 즉 EOS인 경우, 확률값에 -inf을 덮어씀
+        cumulative_prob = self.cumulative_probs[-1].masked_fill_(self.masks[-1], -float('inf'))
+        # 각 단어(output_Size)만큼의 누적확률 값을 계산하기 위해
+        # (beam_size, 1, output_size)의 크기로 늘려줌
+        # 그 후, 입력받은 y_hat과 더해서 최종 누적 확률 값 산출
+        # 하지만 이떄, 맨처음 cumulative_prob가 (0, -inf, -inf)이므로
+        # 처음에는 첫 번째 빔에서만 모든 결과가 나오게 될 것임
+        cumulative_prob = y_hat + cumulative_prob.view(-1, 1, 1).expand(self.beam_size, 1, output_size)
+        # |cumulative_prob| = (beam_size, 1, output_size)
+
+        # cumulative_prob를 (beam_size * output_size,)로 
+        # flatten 해준후 확률이 높은 순으로 정렬
+        # top_indice에는 원래 정렬되기전 index가 유지됨
+        top_log_prob, top_indice = cumulative_prob.view(-1).sort(descending=True)
+        # 그후, TopK개만큼 잘라냄
+        top_log_prob, top_indice = top_log_prob[:self.beam_size], top_indice[:self.beam_size]
+        # |top_log_prob| = (beam_size,)
+        # |top_indice| = (beam_size,)
+        # top_log_prob: 각 단어에 대한 확률 값
+        # top_indice: 각 단어들의 index -> 해당 인덱스를 이용해서
+        # 어느 빔의 어느 단어가 인지를 추적할 수 있음
+       
+
+        # 모든 top_indice를 output_size로 나눈 나머지를 구함으로써
+        # 각 top_indice가 원래 가르키던 word_Index가 튀어나오게 됨
+        self.word_indice += [top_indice.fmod(output_size)]
+        # 모든 top_indice를 output_size로 나눔으로써
+        # 각 top_indice가 원래 가르키던 Beam_index가 나오게 됨
+        # 이로써, 최종적으로 topK에 (어떤 빔)에서 나와서 (어떤 단어)가 선정되었는지 식별
+        self.beam_indice += [top_indice.div(float(output_size)).long()]
+
+        # 이번 스텝에서 구한 누적 확률값을 객체에 갱신
+        self.cumulative_probs += [top_log_prob]
+        # 이번 결과를 보며, EOS가 나온 곳을 mask 처리
+        self.masks += [torch.eq(self.word_indice[-1], data_loader.EOS)] 
+        # 마스크 결과를 바탕으로 done_cnt 캐싱
+        self.done_cnt += self.masks[-1].float().sum()
+
+        # 현재 타임스텝에서 도출된 각종 hidden, cell, h_tilde 값을
+        # 객체에 저장해야 함. -> 이후 get_batch에서 호출될때 사용
+        # 단 이때, topK로 선정된 Beam_index의 hidden, cell, h_tilde만 가지고감
+        for prev_status_name, prev_status in prev_status.items():
+            self.prev_status[prev_status_name] = torch.index_select(
+                prev_status,
+                dim=self.batch_dims[prev_status_name],
+                index=self.beam_indice[-1]
+            ).contiguous()
+
+    def get_n_best(self, n=1, length_penalty=.2):
+        '''
+        이때까지의 Beam Board를 찾아보며,
+        가장 확률 값이 높았던 N개의 문장 추출
+        '''
+        sentences, probs, founds = [], [], []
+        
+        '''
+        mask 여부를 통해, EOS 즉, 온전히 끝난 문장을 탐색
+        찾았다면, 해당 문장의 EOS(끝) 인덱스와 마지막으로 나왔던 beam 인덱스,
+        그리고 그 당시에 누적 확률값을 저장
+        '''
+        for t in range(len(self.word_indice)):  # for each time-step,
+            for b in range(self.beam_size):  # for each beam,
+                if self.masks[t][b] == 1:  # if we had EOS on this time-step and beam,
+                    # Take a record of penaltified log-proability.
+                    probs += [self.cumulative_probs[t][b] * self.get_length_penalty(t, alpha=length_penalty)]
+                    founds += [(t, b)]
+
+        # 만약에, EOS는 아니지만, max_length에 도달해버려 끊겨버린 경우도 수집해옴
+        # Also, collect log-probability from last time-step, for the case of EOS is not shown.
+        for b in range(self.beam_size):
+            if self.cumulative_probs[-1][b] != -float('inf'): # If this beam does not have EOS,
+                if not (len(self.cumulative_probs) - 1, b) in founds:
+                    probs += [self.cumulative_probs[-1][b] * self.get_length_penalty(len(self.cumulative_probs),
+                                                                                     alpha=length_penalty)]
+                    founds += [(t, b)]
+
+        # Sort and take n-best.
+        # 갖고온 문장의 EOS 인덱스를 확률과 묶어서 내림차순 정렬후, N개를 자름
+        sorted_founds_with_probs = sorted(
+            zip(founds, probs),
+            key=itemgetter(1),
+            reverse=True,
+        )[:n]
+        probs = []
+
+        '''
+        정렬된 각 인덱스(EOS)부터 문장을 역으로 내려가며 단어를 수집함
+        이때, 단어가 beam을 계속해서 옮겨다니며 선정했을 것이기에
+        반대로, 자신이 나왔던 beam의 단어를 하나씩 추적하며 내려가야 함
+        '''
+        for (end_index, b), prob in sorted_founds_with_probs:
+            sentence = []
+
+            # Trace from the end.
+            for t in range(end_index, 0, -1):
+                sentence = [self.word_indice[t][b]] + sentence
+                b = self.beam_indice[t][b]
+
+            sentences += [sentence]
+            probs += [prob]
+
+        return sentences, probs
diff --git a/src/11_seq2seq/modules/seq2seq.py b/src/11_seq2seq/modules/seq2seq.py
@@ -3,7 +3,7 @@
 from torch.nn.utils.rnn import pack_padded_sequence as pack
 from torch.nn.utils.rnn import pad_packed_sequence as unpack
 import modules.data_loader as data_loader
-
+from modules.search import SingleBeamSearchBoard
 
 class Encoder(nn.Module):
 
@@ -404,4 +404,160 @@ def search(self, src, is_greedy=True, max_length=255):
         y_hats = torch.cat(y_hats, dim=1)
         indice = torch.cat(indice, dim=1)
 
-        return y_hats, indice
+        return y_hats, indice
+
+    def batch_beam_search(
+        self,
+        src,
+        beam_size=5,
+        max_length=255,
+        n_best=1,
+        length_penalty=.2
+    ):
+        mask, x_length = None, None
+
+        if isinstance(src, tuple):
+            x, x_length = src
+            mask = self.generate_mask(x, x_length)
+            # |mask| = (batch_size, length)
+        else:
+            x = src
+        batch_size = x.size(0)
+
+        emb_src = self.emb_src(x)
+        h_src, h_0_tgt = self.encoder((emb_src, x_length))
+        # |h_src| = (batch_size, length, hidden_size)
+        h_0_tgt = self.merge_encoder_hiddens(h_0_tgt)
+
+        '''
+        initialize 'SingleBeamSearchBoard'
+        각 배치별로, beam_size만큼 페이크 배치를 생성해주는 클래스 초기화
+        hidden_state: 인코더에서 넘어온 히든 스테이트
+        cell_state: 인코더에서 넘어온 셀 스테이트
+        h_t_1_tilde: 이전 스텝의 예측값(input feeding), 
+        처음에는 없으므로 None
+        '''
+        boards = [SingleBeamSearchBoard(
+            h_src.device,
+            {
+                'hidden_state': {
+                    'init_status': h_0_tgt[0][:, i, :].unsqueeze(1),
+                    'batch_dim_index': 1,
+                }, # |hidden_state| = (n_layers, batch_size, hidden_size)
+                'cell_state': {
+                    'init_status': h_0_tgt[1][:, i, :].unsqueeze(1),
+                    'batch_dim_index': 1,
+                }, # |cell_state| = (n_layers, batch_size, hidden_size)
+                'h_t_1_tilde': {
+                    'init_status': None,
+                    'batch_dim_index': 0,
+                }, # |h_t_1_tilde| = (batch_size, 1, hidden_size)
+            },
+            beam_size=beam_size,
+            max_length=max_length,
+        ) for i in range(batch_size)]
+        # 각 보드(batch)들이 예측이 끝났는지 여부
+        # 처음에는 당연히 전부 0으로 이루어짐
+        is_done = [board.is_done() for board in boards]
+
+        length = 0
+        # is_done의 합이 Batch_size를 넘을때까지 반복
+        while sum(is_done) < batch_size and length <= max_length:
+            # current_batch_size = sum(is_done) * beam_size
+
+            # Initialize fabricated variables.
+            # As far as batch-beam-search is running, 
+            # temporary batch-size for fabricated mini-batch is 
+            # 'beam_size'-times bigger than original batch_size.
+            fab_input, fab_hidden, fab_cell, fab_h_t_tilde = [], [], [], []
+            fab_h_src, fab_mask = [], []
+            
+            # 각 input들을 beam_size 만큼 늘려서 가짜 batch_size 생성
+            # input, hidden, cell, h_t_tilde는 이미 보드에서 늘려진 상태
+            # h_src, mask만 그대로 expand 해주면 됨
+            for i, board in enumerate(boards):
+                # Batchify if the inference for the sample is still not finished.
+                if board.is_done() == 0:
+                    # 여기서 현재 타임스텝에 필요한 가짜 batch 데이터 반환
+                    y_hat_i, prev_status = board.get_batch()
+                    hidden_i    = prev_status['hidden_state']
+                    cell_i      = prev_status['cell_state']
+                    h_t_tilde_i = prev_status['h_t_1_tilde']
+
+                    fab_input  += [y_hat_i]
+                    fab_hidden += [hidden_i]
+                    fab_cell   += [cell_i]
+                    fab_h_src  += [h_src[i, :, :]] * beam_size
+                    fab_mask   += [mask[i, :]] * beam_size
+                    if h_t_tilde_i is not None:
+                        fab_h_t_tilde += [h_t_tilde_i]
+                    else:
+                        fab_h_t_tilde = None
+
+            fab_input  = torch.cat(fab_input,  dim=0)
+            fab_hidden = torch.cat(fab_hidden, dim=1)
+            fab_cell   = torch.cat(fab_cell,   dim=1)
+            fab_h_src  = torch.stack(fab_h_src)
+            fab_mask   = torch.stack(fab_mask)
+            if fab_h_t_tilde is not None:
+                fab_h_t_tilde = torch.cat(fab_h_t_tilde, dim=0)
+            # |fab_input|     = (current_batch_size, 1)
+            # |fab_hidden|    = (n_layers, current_batch_size, hidden_size)
+            # |fab_cell|      = (n_layers, current_batch_size, hidden_size)
+            # |fab_h_src|     = (current_batch_size, length, hidden_size)
+            # |fab_mask|      = (current_batch_size, length)
+            # |fab_h_t_tilde| = (current_batch_size, 1, hidden_size)
+
+            emb_t = self.emb_dec(fab_input)
+            # |emb_t| = (current_batch_size, 1, word_vec_size)
+
+            fab_decoder_output, (fab_hidden, fab_cell) = self.decoder(emb_t,
+                                                                      fab_h_t_tilde,
+                                                                      (fab_hidden, fab_cell))
+            # |fab_decoder_output| = (current_batch_size, 1, hidden_size)
+            context_vector = self.attn(fab_h_src, fab_decoder_output, fab_mask)
+            # |context_vector| = (current_batch_size, 1, hidden_size)
+            fab_h_t_tilde = self.tanh(self.concat(torch.cat([fab_decoder_output,
+                                                             context_vector
+                                                             ], dim=-1)))
+            # |fab_h_t_tilde| = (current_batch_size, 1, hidden_size)
+            y_hat = self.generator(fab_h_t_tilde)
+            # |y_hat| = (current_batch_size, 1, output_size)
+
+            # 디코더에서는 그대로 한 batch인듯이 병렬연산을 해준뒤,
+            # 각 board에 다시 beam_size만큼 찢어서 보내줌
+            # fab_hidden[:, begin:end, :] = (n_layers, beam_size, hidden_size)
+            # fab_cell[:, begin:end, :]   = (n_layers, beam_size, hidden_size)
+            # fab_h_t_tilde[begin:end]    = (beam_size, 1, hidden_size)
+            cnt = 0
+            for board in boards:
+                if board.is_done() == 0:
+                    # Decide a range of each sample.
+                    begin = cnt * beam_size
+                    end = begin + beam_size
+
+                    # pick k-best results for each sample.
+                    board.collect_result(
+                        y_hat[begin:end],
+                        {
+                            'hidden_state': fab_hidden[:, begin:end, :],
+                            'cell_state'  : fab_cell[:, begin:end, :],
+                            'h_t_1_tilde' : fab_h_t_tilde[begin:end],
+                        },
+                    )
+                    cnt += 1
+
+            is_done = [board.is_done() for board in boards]
+            length += 1
+
+        # pick n-best hypothesis.
+        batch_sentences, batch_probs = [], []
+
+        # Collect the results.
+        for i, board in enumerate(boards):
+            sentences, probs = board.get_n_best(n_best, length_penalty=length_penalty)
+
+            batch_sentences += [sentences]
+            batch_probs     += [probs]
+
+        return batch_sentences, batch_probs
diff --git a/src/11_seq2seq/modules/seq2seq2.py b/src/11_seq2seq/modules/seq2seq2.py
