Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!7830 modify mass beam_search's floordiv and mod to adapt to 310. 

Merge pull request !7830 from linqingke/mass
This commit is contained in:
mindspore-ci-bot 2020-10-28 09:08:08 +08:00 committed by Gitee
parent d4def00f28 a084532a6e
commit 988e1c1782
1 changed files with 18 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
model_zoo/official/nlp/mass/src/transformer/beam_search.py
View File

@ -224,6 +224,11 @@ class BeamSearchDecoder(nn.Cell):
self.one = Tensor(1, mstype.int32)
self.prob_concat = P.Concat(axis=1)
self.greater_equal = P.GreaterEqual()
self.sub = P.Sub()
self.cast = P.Cast()
self.zeroslike = P.ZerosLike()
def one_step(self, cur_input_ids, enc_states, enc_attention_mask, state_log_probs, state_seq, state_finished,
state_length, entire_log_probs):
"""
@ -261,8 +266,19 @@ class BeamSearchDecoder(nn.Cell):
topk_scores, topk_indices = self.topk(flat_scores, self.beam_width)
# convert to beam and word indices, [batch, beam]
beam_indices = self.floor_div(topk_indices, self.vocab_size_tensor)
word_indices = self.mod(topk_indices, self.vocab_size_tensor)
# beam_indices = self.floor_div(topk_indices, self.vocab_size_tensor)
# word_indices = self.mod(topk_indices, self.vocab_size_tensor)
# ======================================================================
# replace floor_div and mod op, since these two ops only support fp16 on
# Ascend310, which will cause overflow.
temp = topk_indices
beam_indices = self.zeroslike(topk_indices)
for _ in range(self.beam_width - 1):
temp = self.sub(temp, self.vocab_size_tensor)
res = self.cast(self.greater_equal(temp, 0), mstype.int32)
beam_indices = beam_indices + res
word_indices = topk_indices - beam_indices * self.vocab_size_tensor
#======================================================================
current_word_pro = self.gather_nd(
log_probs,