每次為我自己的集束搜索實現獲得不同的 output。 但是當我在代碼中添加延遲時就解決了。 為什么？

Question

我正在嘗試對使用 ONNX io 綁定運行的 T5 model 實施波束搜索。 當我運行波束搜索時，它每次都會產生不同的輸出。 但是當我在令牌生成部分添加延遲時，它每次都會生成相同的 output。 但是延遲不是解決方案，因為我正在從頭開始實施波束搜索。 下面是代碼片段。

batch_size=4
num_beams=4
max_length=15
#Encoder prediction.
enc_out=t5_enc(input_ids=input_ids)
#First time prediction dec ids.
gen_dec_first=torch.zeros((batch_size,1),device="cuda",dtype=torch.long)
#Next sequence decoder ids representing the number of beams.
generated_dec = torch.zeros((batch_size*num_beams,1),device="cuda",dtype=torch.long)

#Scores tensor to accomodate the log probabilities of the outputs chosen.
scores_tensor= torch.zeros((batch_size*num_beams,1),device="cuda",dtype=torch.long)


#Preloop prediction.

#Predict for input_ids of batch_size.
dec_outs=t5_dec(gen_dec_first,enc_out)
#Select the top num_beams size tokens from each prediction.
top_k_ele=torch.topk(dec_outs[:,-1,:],k=num_beams,dim=-1)

#Append them as first prediction.
first_token=top_k_ele.indices.flatten().unsqueeze(0)
first_token=torch.transpose(first_token, 0, 1)
generated_dec=torch.cat((generated_dec,first_token),dim=1)

#Previous time step log probabilities.
prev_prob=top_k_ele.values.flatten().unsqueeze(0)
prev_prob=torch.transpose(prev_prob, 0, 1)
#Append the score first.
scores_tensor=torch.cat((scores_tensor,prev_prob),dim=1)

prev_prob=torch.tile(prev_prob,(1,num_beams))

#Repeat the encoder outputs for num_beams.

# enc_copy=enc_out.detach().clone()
# enc_out=torch.repeat_interleave(enc_out,torch.tensor([4,4,4,4],device="cuda"),dim=0)
enc_out=t5_enc(input_ids=input_ids.repeat_interleave(4, dim=0))


for i in range(max_length):
  dec_outs=t5_dec(generated_dec,enc_out)
  top_k_ele=torch.topk(dec_outs[:,-1,:],k=num_beams,dim=-1)
  c_prob=torch.squeeze(top_k_ele.values,dim=1)
  if i==0:
    f_prob=torch.add(prev_prob,c_prob)
  else:
    prev_prob_t=torch.transpose(prev_prob.unsqueeze(0), 0, 1)
    prev_prob_t=torch.tile(prev_prob_t,(1,num_beams))
    f_prob=torch.add(prev_prob_t,c_prob)
  f_probs_obj=f_prob.max(dim=1)
  f_probs=f_probs_obj.values
  f_indices=top_k_ele.indices[:,f_probs_obj.indices]
  prev_prob=f_probs
  #Append the score.
  s_prob=prev_prob.unsqueeze(0)
  s_prob=torch.transpose(s_prob, 0, 1)
  scores_tensor=torch.cat((scores_tensor,s_prob),dim=1)
  cur_tokens=top_k_ele.indices[:,f_probs_obj.indices]
  cur_tokens=cur_tokens[:,0]
  cur_tokens=cur_tokens.unsqueeze(0)
  cur_tokens=torch.transpose(cur_tokens, 0, 1)
  generated_dec=torch.cat((generated_dec,cur_tokens),dim=1)

  # time.sleep(1) #Having a 1-second delay solves the issue.

output 在貪婪搜索的情況下很好。 我也嘗試設置火炬種子，但由於沒有隨機變量，所以沒有效果。

Answer 1

問題是由於在我的代碼執行時 CUDA 個內核仍在后台運行。 添加 torch.cuda.synchronize() 解決了這個問題。

batch_size=4
num_beams=4
max_length=15
#Encoder prediction.
enc_out=t5_enc(input_ids=input_ids)
#First time prediction dec ids.
gen_dec_first=torch.zeros((batch_size,1),device="cuda",dtype=torch.long)
#Next sequence decoder ids representing the number of beams.
generated_dec = torch.zeros((batch_size*num_beams,1),device="cuda",dtype=torch.long)

#Scores tensor to accomodate the log probabilities of the outputs chosen.
scores_tensor= torch.zeros((batch_size*num_beams,1),device="cuda",dtype=torch.long)


#Preloop prediction.

#Predict for input_ids of batch_size.
dec_outs=t5_dec(gen_dec_first,enc_out)
#Select the top num_beams size tokens from each prediction.
top_k_ele=torch.topk(dec_outs[:,-1,:],k=num_beams,dim=-1)

#Append them as first prediction.
first_token=top_k_ele.indices.flatten().unsqueeze(0)
first_token=torch.transpose(first_token, 0, 1)
generated_dec=torch.cat((generated_dec,first_token),dim=1)

#Previous time step log probabilities.
prev_prob=top_k_ele.values.flatten().unsqueeze(0)
prev_prob=torch.transpose(prev_prob, 0, 1)
#Append the score first.
scores_tensor=torch.cat((scores_tensor,prev_prob),dim=1)

prev_prob=torch.tile(prev_prob,(1,num_beams))

#Repeat the encoder outputs for num_beams.

# enc_copy=enc_out.detach().clone()
# enc_out=torch.repeat_interleave(enc_out,torch.tensor([4,4,4,4],device="cuda"),dim=0)
enc_out=t5_enc(input_ids=input_ids.repeat_interleave(4, dim=0))


for i in range(max_length):
  dec_outs=t5_dec(generated_dec,enc_out)
  top_k_ele=torch.topk(dec_outs[:,-1,:],k=num_beams,dim=-1)
  c_prob=torch.squeeze(top_k_ele.values,dim=1)
  if i==0:
    f_prob=torch.add(prev_prob,c_prob)
  else:
    prev_prob_t=torch.transpose(prev_prob.unsqueeze(0), 0, 1)
    prev_prob_t=torch.tile(prev_prob_t,(1,num_beams))
    f_prob=torch.add(prev_prob_t,c_prob)
  f_probs_obj=f_prob.max(dim=1)
  f_probs=f_probs_obj.values
  f_indices=top_k_ele.indices[:,f_probs_obj.indices]
  prev_prob=f_probs
  #Append the score.
  s_prob=prev_prob.unsqueeze(0)
  s_prob=torch.transpose(s_prob, 0, 1)
  scores_tensor=torch.cat((scores_tensor,s_prob),dim=1)
  cur_tokens=top_k_ele.indices[:,f_probs_obj.indices]
  cur_tokens=cur_tokens[:,0]
  cur_tokens=cur_tokens.unsqueeze(0)
  cur_tokens=torch.transpose(cur_tokens, 0, 1)
  generated_dec=torch.cat((generated_dec,cur_tokens),dim=1)
  torch.cuda.synchronize()

  # time.sleep(1) #Having a 1-second delay solves the issue.