Update beam_search.py from T2T
@@ -547,7 +547,7 @@ def beam_search(symbols_to_logits_fn,
finished_flags, states)
def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
- finished_scores, finished_in_finished, unused_states):
+ finished_scores, unused_finished_in_finished, unused_states):
"""Checking termination condition.
We terminate when we decoded up to decode_length or the lowest scoring item
@@ -559,30 +559,32 @@ def beam_search(symbols_to_logits_fn,
alive_log_probs: probabilities of the beams. [batch_size, beam_size]
finished_scores: scores for each of these sequences.
[batch_size, beam_size]
- finished_in_finished: finished bools for each of these sequences.
- [batch_size, beam_size]
Returns:
Bool.
"""
- if not stop_early:
- return tf.less(i, decode_length)
max_length_penalty = tf.pow(((5. + tf.cast(decode_length, tf.float32)) / 6.), alpha)
# The best possible score of the most likely alive sequence.
lower_bound_alive_scores = alive_log_probs[:, 0] / max_length_penalty
- # Now to compute the lowest score of a finished sequence in finished
- # If the sequence isn't finished, we multiply it's score by 0. since
- # scores are all -ve, taking the min will give us the score of the lowest
- # finished item.
- lowest_score_of_finished_in_finished = tf.reduce_min(
- finished_scores * tf.cast(finished_in_finished, tf.float32), axis=1)
- # If none of the sequences have finished, then the min will be 0 and
- # we have to replace it by -ve INF if it is. The score of any seq in alive
- # will be much higher than -ve INF and the termination condition will not
- # be met.
- lowest_score_of_finished_in_finished += (
- (1. - tf.cast(tf.reduce_any(finished_in_finished, 1), tf.float32)) * -INF)
+ if not stop_early:
+ # by considering the min score (in the top N beams) we ensure that
+ # the decoder will keep decoding until there is at least one beam
+ # (in the top N) that can be improved (w.r.t. the alive beams).
+ # any unfinished beam will have score -INF - thus the min
+ # will always be -INF if there is at least one unfinished beam -
+ # which means the bound_is_met condition cannot be true in this case.
+ lowest_score_of_finished_in_finished = tf.reduce_min(finished_scores)
+ else:
+ # by taking the max score we only care about the the first beam;
+ # as soon as this first beam cannot be beaten from the alive beams
+ # the beam decoder can stop.
+ # similarly to the above, if the top beam is not completed, its
+ # finished_score is -INF, thus it will not activate the
+ # bound_is_met condition. (i.e., decoder will keep going on).
+ # note we need to find the max for every sequence eparately - so, we need
+ # to keep the batch dimension (see axis=1)
+ lowest_score_of_finished_in_finished = tf.reduce_max(finished_scores, axis=1)
bound_is_met = tf.reduce_all(
tf.greater(lowest_score_of_finished_in_finished,
