I have trained two text classification models using GPU on Azure. The models are the following

  1. Bert (ktrain)
  2. Lstm Word2Vec (tensorflow)

Exaples of the code can be found here: nlp

I saved the models into files (.h5) for later use. The files are big e.g. 27,613kb for the lstm and 1.2 gb for bert.

I loaded the models and in a computer where only CPU is available. They both work fine but the model.predict(text) function is super slow predicting the class of the text e.g. on average 1 tweet sized message per second.

Adding GPU on the computer is not an option. I wonder if there is another way to make it run faster? e.g. train the models in a different way (without compromising accuracy) or save the model in a different file format?

  • $\begingroup$ I don't think there is any way to speed up the actual prediction time of the exact same model. The fastest way I can think of is loading the model once, and then predicting in batches. Do you do with torch.no_grad() and set model.eval()? $\endgroup$
    – N. Kiefer
    Oct 15, 2020 at 9:52
  • 1
    $\begingroup$ @N.Kiefer the question specifies that the models are for tensorflow (ktrain is a tf.keras wrapper). $\endgroup$
    – noe
    Oct 15, 2020 at 13:32

1 Answer 1


There are a couple of options:

  • Optimize tensorflow for your specific CPU. Sometimes the official versions of tensorflow are not compiled with support for some instruction sets (e.g. SSE4.1, SSE4.2, AVX, AVX2, FMA). Usually, there is a tensorflow runtime warning message stating so. This prevents some computations to take place in parallel. You can either download a version that is optimized for your CPU (e.g. from intel) or you can compile tensorflow yourself.

  • Prune the model. BERT is a multi-head attention model, and it is possible to remove some of its attention heads while retaining most of the quality. Here's a blog post that shows how to do it for tensorflow.

  • $\begingroup$ thank you for the quick reply. While researching this particular problem I came across TensorFlow Serving. Do you think this could help? tensorflow.org/tfx/tutorials/serving/rest_simple $\endgroup$ Oct 15, 2020 at 15:11
  • 1
    $\begingroup$ Tensorflow serving is just for exposing your model as a web service. The tensorflow runtime is the same, so you will not get better performance just for using tensorflow serving. $\endgroup$
    – noe
    Oct 15, 2020 at 15:14

Your Answer

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

Not the answer you're looking for? Browse other questions tagged or ask your own question.