Steven Reitsma
Jan 21 · 5 min read

Introduction

Google offered us a chance to test their new TPUv2 devices for free on Google Cloud as part of the TensorFlow Research Cloud program. In this blog post we will show some of the experiments we ran to test model accuracy, development speed and runtime. We have found that TPUs have an insanely fast data throughput relative to their cost, but model accuracy sometimes lags behind models trained on more conventional devices such as CPUs and GPUs.

Experiment

We ran all of our experiments on a relatively small convolutional network since we didn’t have access to the TPUs for that long. This definitely skews the results somewhat as distributed learning (in both GPUs and TPUs) works best when using larger networks and datasets. The networks were trained from scratch on Fashion-MNIST. Note that this custom network is not officially optimized for TPUs.

We ran the experiment with the following devices:

  • CPU (no distributed learning)
  • 1 TPUv2 with only 1 core active (no distributed learning)
  • 1 TPUv2 with 8 cores (distributed learning)
  • 1 Tesla V100 GPU (no distributed learning)
  • 8 Tesla V100 GPU’s (distributed learning)

All of the experiments were run on a Google Compute n1-standard-2 machine with 2 CPU cores and 7.5GB of memory, with the exception of the experiment with 8 Tesla V100 GPU’s, where 30GB of memory was given to the machine due to excessive swapping.

To account for any differences in the convergence speed of models (primarily due to data parallelism) we do not train a specific amount of epochs, but instead train until a validation threshold of 0.25 (log loss) is reached. This makes for a fairer comparison.

Below we show a comparison of the training runs between these five device types and show some interesting results. For every run we have tested both tf.train.Adam optimizer (with default settings) and the SGD optimizer with Nesterov momentum (0.9): tf.train.MomentumOptimizer.

CPU baseline

In this graph we can see that on CPU, the network reaches a low validation loss rather quickly. In this case Adam reaches a low point a lot faster than SGD but does overfit in the end. Adam being faster makes sense as we didn’t do any learning rate tuning for the SGD optimizer.

Single GPU: similar results

Again, we see the network converging quickly and that Adam is faster than SGD but does overfit a bit in the end. The most important thing to note is that training a network on a GPU does not give different results compared to CPU.

Data parallelization using 8 GPUs

When we train the network on 8 GPUs using data parallelization, we can see that training requires the same amount of epochs to reach the minimum validation loss. Since data parallelization decreases the number of effective updates per epoch, we increase the learning rate linearly with the batch size, as suggested here. In our case this means increasing the learning rate 8-fold. In terms of wall time, using 8 GPUs is significantly faster than using a single GPU. In larger networks, the speedup would be close to 4–7x. In our smaller network, the speedup is slightly smaller with a 3.7x speedup.

TPU with 8 cores

We now take a look at how the performance of TPUs compares to GPUs.

Comparing TPUs with GPUs:

The SGD optimizer behaves similarly for 8 TPUs as for 8 GPUs, as would be expected. Initially, the Adam optimizer behaved much worse in our TPU experiments because we were using Keras’ Adam optimizer instead of TensorFlow’s tf.train.Adam. After changing our code to use the tf.train.Adam optimizer, Adam performed the same as in GPUs, albeit a bit slower. We also show the results for non-distributed learning for a single TPU core and a single GPU to indicate learning is similar:

Cost comparison

We’ve made a cost comparison to show the cost per training for all five tested devices. The running time of a network is an average of 5 runs, all trained with Adam until a 0.25 validation loss was reached. We show both the figures for regular machines and preemptible machines.

We can see that TPUs are most cost-efficient for this network. We should note that for larger networks these comparisons could differ a lot. For example, both the 8 GPUs and TPUs will see a speedup compared to non-distributed training.

Conclusion

TPUs are fast and cost-efficient, as tested in our smaller example networks. For these smaller networks and datasets we can see that the overhead of parallelization is high, but this is no different from using multiple GPUs.

About the author

Steven Reitsma is a machine learning engineer at BigData Republic, a data science consultancy company in the Netherlands. We hire the best of the best in Data Science and Data Engineering. If you are interested in using machine learning techniques on practical use cases, feel free to contact us at info@bigdatarepublic.nl.

Additional resources

bigdatarepublic

DATA SCIENCE | BIG DATA ENGINEERING | BIG DATA ARCHITECTURES

Thanks to Gerben Oostra and Ruurtjan Pul

Steven Reitsma

Written by

Machine Learning Engineer at BigData Republic

bigdatarepublic

DATA SCIENCE | BIG DATA ENGINEERING | BIG DATA ARCHITECTURES

Welcome to a place where words matter. On Medium, smart voices and original ideas take center stage - with no ads in sight. Watch
Follow all the topics you care about, and we’ll deliver the best stories for you to your homepage and inbox. Explore
Get unlimited access to the best stories on Medium — and support writers while you’re at it. Just $5/month. Upgrade