A Tensor Processing Unit (TPU) is a specialized application-specific chip for deep learning. Google provides TPU access through Google Cloud and Colab as an alternative to GPUs when training speed and cost matter. section 18.2.3
Use TPU training with these practical considerations:
- Expect substantial speed gains: TPU v2 on Colab typically trains about 15× faster than an NVIDIA P100 GPU and averages about 3× better cost-effectiveness. section 18.2.3
- Select a TPU runtime in Colab, or use TPU v3 through v5 on Google Cloud for more serious training runs. section 18.2.3
- Configure distribution before creating the model when using Keras with the JAX backend. section 18.2.3
- Watch input throughput: fast TPU computation can make Google Cloud Storage reads the bottleneck, so cache small datasets in memory. section 18.2.3
- Use very large batches to keep TPU cores busy; small models may require batches exceeding 10,000 samples. section 18.2.3
Starting TPU distribution
With a TPU-enabled notebook and the Keras JAX backend, configure a DataParallel or ModelParallel distribution before creating the model: keras.distribution.set_distribution(distribution). section 18.2.3
model.compile(..., steps_per_execution=8)
Increase work completed per call with step fusing: model.compile(..., steps_per_execution=8). section 18.2.3
