手撕BeamSearch代码，2024年最新花费近一年时间整理的大数据开发核心知识清单

2401_84187563

于 2024-04-08 20:52:01 发布

阅读量711

点赞数 15

分类专栏： 2024年程序员学习文章标签：大数据 pytorch 人工智能

本文链接：https://blog.csdn.net/2401_84187563/article/details/137523465

版权

2024年程序员学习专栏收录该内容

45 篇文章 0 订阅

订阅专栏

):
batch_beam_idx=batch_idx*num_beams+next_index

next_beam_scores[batch_idx, beam_idx] = next_score #当前路径得分
next_beam_tokens[batch_idx, beam_idx] = next_token #当前时刻的token
next_beam_indices[batch_idx, beam_idx] = batch_beam_idx #先前对应的id

beam_idx += 1

return next_beam_scores.view(-1), next_beam_tokens.view(-1), next_beam_indices.view(-1)

beam_scores, beam_next_tokens, beam_idx=process(input_ids,next_token_scores,next_tokens,next_indices)

更新输入，找到对应的beam_idx, 选择的tokens, 拼接为新的输入 #(batch*beam,seq_len)

input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
cur_len = cur_len + 1
#输出
return input_ids,beam_scores

if name == ‘main’:
input_ids=torch.randint(0,100,size=(3,1))
print(input_ids)
input_ids,beam_scores=beam_search(input_ids,max_length=10,num_beams=3)
print(input_ids)

参考：transformers generate实现。

transformer generate() 解读

@torch.no_grad()
def generate( #模型入口
self,
inputs: Optional[torch.Tensor] = None,
generation_config: Optional[GenerationConfig] = None,
logits_processor: Optional[LogitsProcessorList] = None,
stopping_criteria: Optional[StoppingCriteriaList] = None,
prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
synced_gpus: Optional[bool] = None,
assistant_model: Optional[“PreTrainedModel”] = None,
streamer: Optional[“BaseStreamer”] = None,
negative_prompt_ids: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
**kwargs,
) -> Union[GenerateOutput, torch.LongTensor]:

10. go into different generation modes

根据不同的生产模型进行解码生产

if generation_mode == GenerationMode.ASSISTED_GENERATION:
…
#以beam search 为例子
elif generation_mode == GenerationMode.BEAM_SEARCH: #beam search 算法

11. prepare beam search scorer #参数初始化

beam_scorer = BeamSearchScorer(
batch_size=batch_size,
num_beams=generation_config.num_beams,
device=inputs_tensor.device,
length_penalty=generation_config.length_penalty,
do_early_stopping=generation_config.early_stopping,
num_beam_hyps_to_keep=generation_config.num_return_sequences,
max_length=generation_config.max_length,
)
#将输入进行扩展

12. interleave input_ids with `num_beams` additional sequences per batch

input_ids, model_kwargs = self._expand_inputs_for_generation(
input_ids=input_ids,
expand_size=generation_config.num_beams,
is_encoder_decoder=self.config.is_encoder_decoder,
**model_kwargs,
)

13. run beam search 核心，beam search 算法解码

result = self.beam_search(
input_ids,
beam_scorer,
logits_processor=prepared_logits_processor,
stopping_criteria=prepared_stopping_criteria,
pad_token_id=generation_config.pad_token_id,
eos_token_id=generation_config.eos_token_id,
output_scores=generation_config.output_scores,
output_logits=generation_config.output_logits,
return_dict_in_generate=generation_config.return_dict_in_generate,
synced_gpus=synced_gpus,
sequential=generation_config.low_memory,
**model_kwargs,
)

def beam_search(
self, input_ids, encoder_output, attention_mask, num_beams, max_length, pad_token_id: int, eos_token_id: int
):
batch_size = self.beam_scorer.batch_size #扩展前batch size

num_beams = self.beam_scorer.num_beams
batch_beam_size, cur_len = input_ids.shape #扩展后batch

assert (
num_beams * batch_size == batch_beam_size
), f"Batch dimension of input_ids should be {num_beams * batch_size}, but is {batch_beam_size}."

beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
beam_scores[:, 1:] = -1e9
beam_scores = beam_scores.view((batch_size * num_beams,))
next_tokens = torch.zeros((batch_size, num_beams), dtype=torch.long, device=input_ids.device)
next_indices = torch.zeros((batch_size, num_beams), dtype=torch.long, device=input_ids.device)

past: List[torch.Tensor] = []
while cur_len < max_length:
#生成相应
logits, past = self._decoder_forward(input_ids, encoder_output, attention_mask, past) #迭代输出
next_token_logits = logits[:, -1, :] #当前时刻输出

adjust tokens for Bart, e.g. cur_len=1 与 max_length 输出调整

next_token_logits = self.adjust_logits_during_generation(
next_token_logits, cur_len=cur_len, max_length=max_length
)
#归一化
next_token_scores = F.log_softmax(next_token_logits, dim=-1) # (batch_size * num_beams, vocab_size) #归一化