I would like to train an LSTM-based variational autoencoder on a large dataset (37 million sentences). However, I have calculated that my training speed as of now is too slow (on Google Colab). I am using a GPU provided by Google called A100-SXM4-40GB, and my framework of choice is Pytorch. I am already using automatic mixed precision, which sped up my code by about x2 (https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html). As a reminder, here is a stock image of a variational autoencoder:VAE

With my current training speed, I get through about 130 million training examples/sentences in 24 hours. My vocabulary size is about 85,000, the number of parameters in the VAE is 17.1 million (mostly high because of the embedding layer), my encoder has 100 and my decoder 600 hidden neurons. My batch size is 64, and I am using the Adam optimizer. I have plotted the time usage of different parts of my code after 3000 batches: Time usage of VAE in percent. Following, you can see a more detailed breakdown of my model times: Percentage of time used by different parts of the VAE

What advice can you give me to speed up my model? For instance, I also have access to a TPU, but I have never seen a clear breakdown on GPU vs TPU performance (and what role batch size plays). Can I use parallel computing, and is this possible on Colab? I am training for 5 epochs, leading to a total estimated training time of about 34 hours.

With regards to the detailed model training breakdown, please note:

  • I am already using pack_padded_sequence for the encoder
  • The decoder is an LSTM cell (less optimised). I need to pass the last encoder hidden state (after reparametrization) to the decoder at each time step, together with the previous output of the cell state for the LSTM. I believe this is following 4

EDIT: I updated my total run time (made a miscalculation)

  • $\begingroup$ Is there any reason the decoder would take so much longer than the encoder? Is the loss computation featured in the 'decoding' part? $\endgroup$ Commented Jan 18, 2023 at 15:18
  • $\begingroup$ The breakdown of the model times (second chart) only refers to the time taken up by the model (in the first chart). So no, the loss computation is not included there. And see my last point why the decoder is taking longer than the encoder (LSTM cell in the decoder out of necessity). I need to do three things in the decoder: 1. keep track of the sentence length (so it does not go beyond a hardcoded max), 2. pass the last encoder hidden state, 3. pass the most recent output of the decoder as its current hidden state. $\endgroup$ Commented Jan 18, 2023 at 15:32
  • $\begingroup$ This is only possible with LSTM cell, which is not optimised to the same degree as a normal LSTM in Pytorch. $\endgroup$ Commented Jan 18, 2023 at 15:37
  • $\begingroup$ Well, as most of the time is spend on your decoder, please provide elaborately how we could help speed up the decoder. Because the conclusion from your own analysis is pretty simply "make the decoding faster" $\endgroup$ Commented Jan 18, 2023 at 21:39
  • $\begingroup$ That is one way, but I think what I am looking for is an analysis of TPU vs GPU comparisons for different batch sizes of RNNs. Because, honestly, I believe my code is fine. I will open a new question for that. You already mentioned in another answer that increasing the batch size is an easy way to increase speed, which I am doing. Thank you for your help! $\endgroup$ Commented Jan 19, 2023 at 19:04


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.