How to speed up the RNN in the for loop

nlp
1

Hi, I am trying to use GRU to predict tokens. However, the computation speed is very slow in for loop.

My code is as below:

decoded_words, word_distributions = [], []

for _ in range(max_utter_len):
    word_distribution, decoder_h = self.decoder(
        decoder_inputs, contexts, decoder_h)

    word_distribution = word_distribution.squeeze(1)
    # [batch size, voc size]

    top1 = word_distribution.topk(1)[1]
    # [batch size, 1]

    decoded_words.append(top1.squeeze(-1))
    word_distributions.append(word_distribution)

    decoder_inputs = top1

decoded_words = torch.stack(decoded_words, dim=-1)
word_distributions = torch.stack(word_distributions, dim=1)

It takes 17 seconds for every 100 iterations.

However, I use the “teacher forcing” in the training, and it only takes 7 seconds for every 100 iterations.
The code is as below:

word_distributions, decoder_init_hidden = self.decoder(
            inputs, context, decoder_init_hidden)

How to speed up RNN prediction?
Thank you.

2

I am guessing this is only for the inference part. The training speed for both the first case and the second case is still pretty similar?

Remember teacher forcing can be used only during training and not during inference.

What you are doing in the first case for inference looks fine to me. If you don’t want to break on <EOS> token, you can prepare a list of the required size (max_utter_len) and then fill it using decoded_words[i] = top1.squeeze(-1) and word_distributions[i] = word_distribution. This will give marginal improvement.

3

They differ by 10 seconds.

I know it. I only use teacher forcing in my training.

Appreciate it. I will try it.