[코드 공유] Beam Search 구현

[shjas94]

def forward(
        self, src, text, is_train=True, batch_max_length=50, teacher_forcing_ratio=1.0, beam_size=4
    ):
        random_v = random.random()
        if is_train and random_v < teacher_forcing_ratio:  # train with teacher forcing
            # teacher forcing시에
            tgt = self.text_embedding(text)
            tgt = self.pos_encoder(tgt)
            tgt_mask = self.pad_mask(text) | self.order_mask(text.size(1))
            for layer in self.attention_layers:
                tgt = layer(tgt, None, src, tgt_mask)
            out = self.generator(tgt)

        elif is_train and random_v >= teacher_forcing_ratio:  # train without teacher forcing
            out = []
            num_steps = batch_max_length - 1
            target = torch.LongTensor(src.size(0)).fill_(
                self.st_id).to(device)  # [START] token
            features = [None] * self.layer_num

            for t in range(num_steps):
                target = target.unsqueeze(1)
                tgt = self.text_embedding(target)
                tgt = self.pos_encoder(tgt, point=t)
                tgt_mask = self.order_mask(t + 1)
                tgt_mask = tgt_mask[:, -1].unsqueeze(1)  # [1, (l+1)]
                for l, layer in enumerate(self.attention_layers):
                    tgt = layer(tgt, features[l], src, tgt_mask)
                    features[l] = (
                        tgt if features[l] == None else torch.cat(
                            [features[l], tgt], 1)
                    )

                _out = self.generator(tgt)  # [b, 1, c]
                target = torch.argmax(_out[:, -1:, :], dim=-1)  # [b, 1]
                target = target.squeeze()   # [b]
                out.append(_out)

            # [b, max length, 1, class length]
            out = torch.stack(out, dim=1).to(device)
            out = out.squeeze(2)    # [b, max length, class length]

        elif not is_train:  # inference or validation (beam search)
            num_steps = batch_max_length - 1
            temp_tar = [
                [torch.LongTensor(src.size(0)).fill_(self.st_id).unsqueeze(1).to(device), torch.LongTensor(src.size(0)).fill_(0).to(device), [None] * self.layer_num, []]]

            k = beam_size
            batch = src.shape[0]
            # 첫 번째 스텝에 input으로 넣어줄 값 임시 저장
            # 스텝 돌면서 append
            for t in range(num_steps):
                new_tar = []
                # 한 스텝마다 뽑혀져 나오는 beam size ~ beam size^2 만큼의 output들 임시 저장
                # [[seq, prob, feature, output], ...]
                # seq : [batch, seq_len], prob : [batch, prob], feature : [batch, ]
                for i, tar in enumerate(temp_tar):
                    target = tar[0][:, -1].unsqueeze(1).to(device)
                    pre_prob = tar[1]
                    pre_feature = tar[2]  # [feature] * layer
                    temp_out = tar[3]
                    tgt = self.text_embedding(target)
                    tgt = self.pos_encoder(tgt, point=t)
                    tgt_mask = self.order_mask(t + 1)
                    tgt_mask = tgt_mask[:, -1].unsqueeze(1)  # [1, (l+1)]

                    for l, layer in enumerate(self.attention_layers):
                        tgt = layer(tgt, pre_feature[l], src, tgt_mask)
                        pre_feature[l] = (
                            tgt if pre_feature[l] == None else torch.cat(
                                [pre_feature[l], tgt], 1)
                        )  # [feature[0], feature[0]+tgt_first, feature[0]+tgt_first+tgt_second, ... ]

                    _out = self.generator(tgt)
                    temp_out.append(_out)  # 일단 append는 정상적으로 동작
                    prob = torch.topk(
                        _out[:, -1:, :], k=beam_size, dim=-1)[0].squeeze().transpose(0, 1)
                    idx = torch.topk(
                        _out[:, -1:, :], k=beam_size, dim=-1)[1].squeeze().transpose(0, 1)

                    for j in range(len(idx)):
                        if t == 0:
                            new_tar.append([torch.stack([tar[0].squeeze().to(device), idx[j].to(
                                device)], dim=1), (prob[j] + pre_prob)/2, pre_feature, temp_out])

                        else:
                            tar_bf_stack = tar[0].transpose(0, 1)
                            idx_bf_stack = idx[j].unsqueeze(1).transpose(0, 1)
                            stack = torch.cat(
                                (tar_bf_stack, idx_bf_stack), dim=0).transpose(0, 1)
                            new_tar.append(
                                [stack, (prob[j] + pre_prob)/2, pre_feature, temp_out])

                stacked_prob = torch.stack(
                    [new_tar[j][1] for j in range(len(new_tar))], dim=-1)

                stacked_seq = torch.stack([new_tar[j][0]
                                          for j in range(len(new_tar))], dim=-1)
                stacked_feature = []
                feature_len = len(new_tar[0][2])
                for f in range(feature_len):
                    temp_feature = []
                    for t_idx in range(len(new_tar)):
                        temp_feature.append(new_tar[t_idx][2][f])
                    stacked_feature.append(torch.stack(temp_feature, dim=-1))
                stacked_out = []
                out_len = len(new_tar[0][3])
                for o in range(out_len):
                    temp_out = []
                    for t_idx in range(len(new_tar)):
                        temp_out.append(new_tar[t_idx][3][o])
                    stacked_out.append(torch.stack(temp_out, dim=-1))

                _, topk_indices = torch.topk(stacked_prob, dim=-1, k=beam_size)
                new_temp_tar = []
                for i in range(beam_size):
                    temp_prob = list()
                    temp_seq = list()
                    temp_feature = [[] for _ in range(len(stacked_feature))]
                    temp_out_beam = [[] for _ in range(len(stacked_out))]
                    for j in range(batch):  # 배치 사이즈
                        temp_prob.append(
                            stacked_prob[j, topk_indices[j, i]])
                        temp_seq.append(
                            stacked_seq[j, :, topk_indices[j, i]])
                        for k in range(len(stacked_feature)):
                            temp_feature[k].append(
                                stacked_feature[k][j, :, :, topk_indices[j, i]])
                        for k in range(len(stacked_out)):
                            temp_out_beam[k].append(
                                stacked_out[k][j, :, :, topk_indices[j, i]])
                    temp_seq_stack = torch.stack(temp_seq, axis=0)
                    temp_prob_stack = torch.stack(temp_prob, axis=0)
                    temp_feature_stack = [torch.stack(
                        temp_feature[j], axis=0) for j in range(len(temp_feature))]
                    temp_out_stack = [torch.stack(
                        temp_out_beam[j], axis=0) for j in range(len(temp_out_beam))]
                    new_temp_tar.append(
                        [temp_seq_stack, temp_prob_stack, temp_feature_stack, temp_out_stack])

                temp_tar = new_temp_tar  # 이렇게 되면 리스트에는 자동적으로 확률값 기준으로 내림차순으로 정렬되어 있음
            out = temp_tar[0][-1]

            out = torch.stack(out, dim=1).to(device)
            out = out.squeeze(2)
        return out

빔서치 구현입니다. TransformerDecoder class의 forward함수만 요걸로 교체해주시면 됩니다. 혹시 제출 파일 있는분들은 적용해서 제출하고 결과 공유해주시면 감사하겠습니다.

[jo-member]

def forward(
            self, src, text, is_train=True, batch_max_length=50, teacher_forcing_ratio=1.0
    ):
        if is_train and random.random() < teacher_forcing_ratio:
            # teacher forcing시에
            tgt = self.text_embedding(text)
            tgt = self.pos_encoder(tgt)
            tgt_mask = self.pad_mask(text) | self.order_mask(text.size(1))
            for layer in self.attention_layers:
                tgt = layer(tgt, None, src, tgt_mask)
            out = self.generator(tgt)
        else:
            print(src.shape)
            num_steps = batch_max_length - 1
            temp_tar = [
                [torch.LongTensor(src.size(0)).fill_(self.st_id).unsqueeze(1).to(device), torch.LongTensor(src.size(0)).fill_(0).to(device), [None] * self.layer_num, []]]
            k = self.k
            for t in range(num_steps):
                new_tar = []
                for i, tar in enumerate(temp_tar):
                    target = tar[0][:,-1].unsqueeze(1).to(device)

                    pre_prob = tar[1]
                    pre_feature = tar[2]
                    temp_out = tar[3]
                    tgt = self.text_embedding(target)
                    tgt = self.pos_encoder(tgt, point=t)
                    tgt_mask = self.order_mask(t + 1)
                    tgt_mask = tgt_mask[:, -1].unsqueeze(1)  # [1, (l+1)]
                    for l, layer in enumerate(self.attention_layers):
                        tgt = layer(tgt, pre_feature[l], src, tgt_mask)
                        pre_feature[l] = (
                            tgt if pre_feature[l] == None else torch.cat([pre_feature[l], tgt], 1)
                        )
                    _out = self.generator(tgt)
                    temp_out.append(_out)
                    prob = torch.topk(_out[:, -1:, :], k=k, dim=-1)[0].squeeze().transpose(0,1)
                    idx = torch.topk(_out[:, -1:, :], k=k, dim=-1)[1].squeeze().transpose(0,1)
                    for i in range(len(idx)):
                        new_tar.append([torch.stack([tar[0].squeeze().to(device), idx[i].to(device)],dim=1), (prob[i] + pre_prob)/2, pre_feature, temp_out])
                # 이제 new_tar안에는 [seq,확률] 들이 저장되어있다. 확률값들을 비교하여 2개만 남긴다음에 나머지 삭제
                sorted_tar = sorted(new_tar, key=lambda x: x[1],reverse=True)
                temp_tar = [sorted_tar[:k]]
            out = sorted(temp_tar,key=lambda x:x[1],reverse=True)[0][-1] # [b, max length, 1, class length]
            out = torch.stack(out,dim=1).to(device)
            out = out.squeeze(2)
        return out

batch를 해결하지 않은 초기 prototype