Supported reference models

Cloud TPU provides the following set of reference models that are optimized for fast and accurate training.

The following table displays the major and minor releases of TensorFlow and PyTorch that are currently supported by Cloud TPU. TensorFlow release numbering has changed with release 2.5.0. Going forward, major TensorFlow release numbers end with '0' and all patch release numbers end with numbers greater than '0'. For example, TF 2.5.0 is a major release and TF 2.4.1 is a minor release. In order to run the latest supported version, check to see if there are any patch releases to the major release. If so, you can run the latest supported patch release rather than the major release.
Framework Major version Model category Reference models Supported versions
TensorFlow 2.x Image classification ResNet-2.x, MNIST-2.x, EfficientNet-2.x 2.1, 2.2, 2.3, 2.4, 2.5.0, 2.6.0
Language modeling Transformer-2.x, BERT-2.x 2.1, 2.2, 2.3, 2.4, 2.5.0, 2.6.0
Language modeling XlNet-2.x 2.4, 2.5.0
Object detection RetinaNet-2.x 2.1, 2.2, 2.3, 2.4, 2.5.0, 2.6.0
Image segmentation Mask-RCNN-2.x 2.2, 2.3, 2.4, 2.5.0, 2.6.0
Image segmentation ShapeMask-2.x 2.3, 2.4, 2.5.0, 2.6.0
Recommendation systems DLRM-2.x, DCN-2.x 2.5.0, 2.6.0
Recommendation systems NCF-2.x 2.3, 2.4, 2.5.0, 2.6.0
1.x Image classification ResNet, AmoebaNet, MNasNet, MNIST, EfficientNet 1.15
Language modeling Transformer, BERT 1.15
Object detection RetinaNet 1.15
PyTorch 1.x Image classification ResNet-PyTorch 1.10
Language modeling FairSeq Transformer, FairSeq RoBERTa Wav2Vec2 1.10
Speech recognition Wav2Vec2 1.10
Recommendation systems DLRM 1.10
JAX/FLAX latest Image classification ResNet50 latest