Commit ef179415 authored by Myle Ott's avatar Myle Ott
Browse files

Faster generation when using a single model (rather than ensemble)

parent bfcc6ec7

fairseq/sequence_generator.py

+17 −12
Changes for fairseq/sequence_generator.py: 17 added lines, 12 removed lines.
Original line number Diff line number Diff line
@@ -493,25 +493,18 @@ class SequenceGenerator(object): 
        return finalized



    def _decode(self, tokens, encoder_outs, incremental_states):

        if len(self.models) == 1:

            return self._decode_one(tokens, self.models[0], encoder_outs[0], incremental_states, log_probs=True)



        avg_probs = None

        avg_attn = None



        for model, encoder_out in zip(self.models, encoder_outs):

            with torch.no_grad():

                if incremental_states[model] is not None:

                    decoder_out = list(model.decoder(tokens, encoder_out, incremental_states[model]))

                else:

                    decoder_out = list(model.decoder(tokens, encoder_out))

                decoder_out[0] = decoder_out[0][:, -1, :]

                attn = decoder_out[1]



            probs = model.get_normalized_probs(decoder_out, log_probs=False).data

            probs, attn = self._decode_one(tokens, model, encoder_out, incremental_states, log_probs=False)

            if avg_probs is None:

                avg_probs = probs

            else:

                avg_probs.add_(probs)

            if attn is not None:

                attn = attn[:, -1, :].data

                if avg_attn is None:

                    avg_attn = attn

                else:
@@ -520,5 +513,17 @@ class SequenceGenerator(object): 
        avg_probs.log_()

        if avg_attn is not None:

            avg_attn.div_(len(self.models))



        return avg_probs, avg_attn



    def _decode_one(self, tokens, model, encoder_out, incremental_states, log_probs):

        with torch.no_grad():

            if incremental_states[model] is not None:

                decoder_out = list(model.decoder(tokens, encoder_out, incremental_states[model]))

            else:

                decoder_out = list(model.decoder(tokens, encoder_out))

            decoder_out[0] = decoder_out[0][:, -1, :]

            attn = decoder_out[1]

            if attn is not None:

                attn = attn[:, -1, :]

        probs = model.get_normalized_probs(decoder_out, log_probs=log_probs)

        return probs, attn